transforms中使用自定义函数

有篇博客写的很清晰，详见：https://blog.csdn.net/qq_32531505/article/details/106300540
使用如下的方法1添加自定义的transform增强

class  方法名(object):
    def __init__(self, x1,x2...):
        self.x1 = x1   #此处是初始化的时候把所需要的的参数加到方法中
        self.x2 = x2  
      
    def __call__(self, img,x1,x2,...):
        写详细的方法内容
        return img

功能：在训练时以0.5的概率进行人脸patch shuffle, 测试时只进行to_tensor到0，1之间，亲测此增强对模型性能有较大提升

代码如下：

class face_destruction(object):
    def __init__(self, ratio = 0.5):
        self.ratio = ratio
        
    def __call__(self, image):
        if np.random.random() > self.ratio:
            width, height = image.size
            item_width = int(width / 2)
            box_list = []
            # (left, upper, right, lower)
            for i in range(0, 2):
                for j in range(0, 2):
                    #print((i*item_width,j*item_width,(i+1)*item_width,(j+1)*item_width))
                    box = (j*item_width, i*item_width, (j+1)*item_width, (i+1)*item_width)
                    box_list.append(box)
        
            image_list = [np.array(image.crop(box)).reshape(56*56*3, 1) for box in box_list]
            
            #image_list = [image.crop(box) for box in box_list]
            random.shuffle(image_list)
            image_list = [i_list.reshape(56, 56, 3) for i_list in image_list]
        
            shuffle_image = Image.new('RGB', (112, 112), color='white')
        
            for p in range(1, 3):
                for q in range(1, 3):
                    fimage = image_list[2 * (p - 1) + q - 1]
                    shuffle_image.paste(Image.fromarray(fimage), (56*(q-1), 56*(p-1), 56*(q-1)+56, 56*p))
        else:
            shuffle_image = image
    
        return shuffle_image

调用时：

from datasets.FASDataset import face_destruction
train_transform = transforms.Compose([
    transforms.RandomRotation(cfg['dataset']['augmentation']['rotation_range']),
    transforms.RandomHorizontalFlip(),
    Cutout(0.3),
    face_destruction(),
    transforms.ToTensor(),
])

val_transform = transforms.Compose([
    transforms.ToTensor(),
])