如何在 Pytorch 中使用 torchvision.transforms 对分割任务进行数据增强？

Posted 2023-03-12

【中文标题】如何在 Pytorch 中使用 torchvision.transforms 对分割任务进行数据增强？

【英文标题】：How to use torchvision.transforms for data augmentation of segmentation task in Pytorch?

【发布时间】：2020-02-01 12:46:33

【问题描述】：

我对 PyTorch 中执行的数据增强有点困惑。

因为我们在处理分割任务，我们需要数据和掩码来进行相同的数据增强，但其中一些是随机的，例如随机旋转。

Keras 提供了random seed 保证 data 和 mask 做同样的操作，如下代码所示：

    data_gen_args = dict(featurewise_center=True,
                         featurewise_std_normalization=True,
                         rotation_range=25,
                         horizontal_flip=True,
                         vertical_flip=True)


    image_datagen = ImageDataGenerator(**data_gen_args)
    mask_datagen = ImageDataGenerator(**data_gen_args)

    seed = 1
    image_generator = image_datagen.flow(train_data, seed=seed, batch_size=1)
    mask_generator = mask_datagen.flow(train_label, seed=seed, batch_size=1)

    train_generator = zip(image_generator, mask_generator)

我在Pytorch官方文档中没有找到类似的描述，所以不知道如何保证数据和掩码可以同步处理。

Pytorch 确实提供了这样的功能，但我想将其应用于自定义 Dataloader。

例如：

def __getitem__(self, index):
    img = np.zeros((self.im_ht, self.im_wd, channel_size))
    mask = np.zeros((self.im_ht, self.im_wd, channel_size))

    temp_img = np.load(Image_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_label = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')

    for i in range(channel_size):
        img[:,:,i] = temp_img[self.count[index] + i]
        mask[:,:,i] = temp_label[self.count[index] + i]

    if self.transforms:
        img = np.uint8(img)
        mask = np.uint8(mask)
        img = self.transforms(img)
        mask = self.transforms(mask)

    return img, mask

这种情况下img和mask会分开变换，因为一些随机旋转等操作是随机的，所以mask和image的对应关系可能会改变。换句话说，图像可能已经旋转，但蒙版没有这样做。

编辑 1

我用augmentations.py中的方法，但是报错：

Traceback (most recent call last):
  File "test_transform.py", line 87, in <module>
    for batch_idx, image, mask in enumerate(train_loader):
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in __next__
    batch = self.collate_fn([self.dataset[i] for i in indices])
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 314, in <listcomp>
    batch = self.collate_fn([self.dataset[i] for i in indices])
  File "/home/dirk/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 103, in __getitem__
    return self.dataset[self.indices[idx]]
  File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/data.py", line 164, in __getitem__
    img, mask = self.transforms(img, mask)
  File "/home/dirk/home/data/dirk/segmentation_unet_pytorch/augmentations.py", line 17, in __call__
    img, mask = a(img, mask)
TypeError: __call__() takes 2 positional arguments but 3 were given

这是我__getitem__()的代码：

data_transforms = 
    'train': Compose([
        RandomHorizontallyFlip(),
        RandomRotate(degree=25),
        transforms.ToTensor()
    ]),


train_set = DatasetUnetForTestTransform(fold=args.fold, random_index=args.random_index,transforms=data_transforms['train'])

# __getitem__ in class DatasetUnetForTestTransform
def __getitem__(self, index):
    img = np.zeros((self.im_ht, self.im_wd, channel_size))
    mask = np.zeros((self.im_ht, self.im_wd, channel_size))
    temp_img = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_label = np.load(Label_path + ':0>4'.format(self.patient_index[index]) + '.npy')
    temp_img, temp_label = crop_data_label_from_0(temp_img, temp_label)
    for i in range(channel_size):
        img[:,:,i] = temp_img[self.count[index] + i]
        mask[:,:,i] = temp_label[self.count[index] + i]

    if self.transforms:
        img = T.ToPILImage()(np.uint8(img))
        mask = T.ToPILImage()(np.uint8(mask))
        img, mask = self.transforms(img, mask)

    img = T.ToTensor()(img).copy()
    mask = T.ToTensor()(mask).copy()
    return img, mask

编辑 2

我发现ToTensor之后，相同标签之间的骰子变成255而不是1，如何解决？

# Dice computation
def DSC_computation(label, pred):
    pred_sum = pred.sum()
    label_sum = label.sum()
    inter_sum = np.logical_and(pred, label).sum()
    return 2 * float(inter_sum) / (pred_sum + label_sum)

随时询问是否需要更多代码来解释问题。

【参考方案1】：

需要像RandomCrop 这样的输入参数的转换有一个get_param 方法，它将返回该特定转换的参数。然后可以使用转换的功能接口将其应用于图像和蒙版：

from torchvision import transforms
import torchvision.transforms.functional as F

i, j, h, w = transforms.RandomCrop.get_params(input, (100, 100))
input = F.crop(input, i, j, h, w)
target = F.crop(target, i, j, h, w)

此处提供示例： https://github.com/pytorch/vision/releases/tag/v0.2.0

此处提供 VOC 和 COCO 的完整示例： https://github.com/pytorch/vision/blob/master/references/segmentation/transforms.py https://github.com/pytorch/vision/blob/master/references/segmentation/train.py

关于错误，

ToTensor() 未被覆盖以处理额外的掩码参数，因此它不能在data_transforms 中。此外，__getitem__ 在返回之前对img 和mask 执行ToTensor。

data_transforms = 
    'train': Compose([
        RandomHorizontallyFlip(),
        RandomRotate(degree=25),
        #transforms.ToTensor()  => remove this line
    ]),

【讨论】：

感谢您的回答，我使用该方法实现了transform，但是在问题描述中添加了一个错误。

用修复更新了答案。

【参考方案2】：

torchvision也提供了类似的功能[document]。

这是一个简单的例子，

import torchvision
from torchvision import transforms

trans = transforms.Compose([transforms.CenterCrop((178, 178)),
                                    transforms.Resize(128),
                                    transforms.RandomRotation(20),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
dset = torchvision.datasets.MNIST(data_root, transforms=trans)

---

编辑

自定义您自己的 CelebA 数据集时的一个简短示例。请注意，要应用转换，您需要在__getitem__ 中调用transform 列表。

class CelebADataset(Dataset):
    def __init__(self, root, transforms=None, num=None):
        super(CelebADataset, self).__init__()

        self.img_root = os.path.join(root, 'img_align_celeba')
        self.attr_root = os.path.join(root, 'Anno/list_attr_celeba.txt')
        self.transforms = transforms

        df = pd.read_csv(self.attr_root, sep='\s+', header=1, index_col=0)
        #print(df.columns.tolist())
        if num is None:
            self.labels = df.values
            self.img_name = df.index.values
        else:
            self.labels = df.values[:num]
            self.img_name = df.index.values[:num]

    def __getitem__(self, index):
        img = Image.open(os.path.join(self.img_root, self.img_name[index]))
        # only use blond_hair, eyeglass, male, smile
        indices = [9, 15, 20, 31]
        label = np.take(self.labels[index], indices)
        label[label==-1] = 0

        if self.transforms is not None:
            img = self.transforms(img)

        return np.asarray(img), label

    def __len__(self):
        return len(self.labels)

---

编辑 2

乍一看，我可能会错过一些东西。您问题的重点是如何将“相同”的数据预处理应用于 img 和标签。据我了解，没有可用的 Pytorch 内置函数。所以，我之前所做的就是自己实现增强。

class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, img, mask):
        rotate_degree = random.random() * 2 * self.degree - self.degree
        return img.rotate(rotate_degree, Image.BILINEAR), 
                           mask.rotate(rotate_degree, Image.NEAREST)

请注意，输入应为 PIL 格式。请参阅this 了解更多信息。

【讨论】：

我重新编辑了我的问题。您的示例用于分类任务。如果用于分割任务如何修改它？具体来说，我们需要同时旋转图片和标签，那么这里应该如何保证它们的对应呢？如果使用img, label = self.transforms(img), self.transforms(label)，变换中的随机旋转会破坏图像和标签之间的对应关系。换句话说，图像可能已经旋转，但蒙版没有这样做。

@Dirk Li，感谢您的评论。我已经添加了我之前的处理方式。

我使用了augmentations.py中的方法，但是出现了错误。并且我在问题描述中添加了错误描述。

我的错。请在augmentation.py 中使用Compose。我已经相应地更新了我的答案。

感谢您的回答，我现在知道问题的原因了。但是还有一个地方让我很困惑。我使用labe1, label2 = self.transforms(label,label) 来测试transform 是否改变了数据。我用label1 和label2 计算骰子，结果是255（以前是1）。我知道ToTensor会把数据除以255，但是我还是不明白Dice为什么变成255，怎么解决。顺便说一句，增强中的 random_rotation 会将 Dice 减少到 244。这是否意味着它改变了对应关系？将骰子计算添加到问题描述中。

【参考方案3】：

另一个想法是沿着通道维度堆叠图像和蒙版，然后将它们一起转换。显然，这只适用于几何类型的转换，您需要对两者使用相同的 dtype。我使用这样的东西：

# Apply these to image and mask
affine_transforms = transforms.Compose([
    transforms.RandomAffine(degrees=180),
    ...
])

# Apply these to image only
image_transforms = transforms.Compose([
    transforms.GaussianBlur(),
    ...
])

# Loader...
def __getitem__(self, index: int):
    # Get the image and mask, here shape=(HxW) for both
    image = self.images[index]
    mask = self.masks[index]

    # Stack the image and mask together so they get the same geometric transformations
    stacked = torch.cat([image, mask], dim=0)  # shape=(2xHxW)
    stacked = self.affine_transforms(stacked)

    # Split them back up again
    image, mask = torch.chunk(stacked, chunks=2, dim=0)

    # Image transforms are only applied to the image
    image = self.image_transforms(image)

    return image, mask