语义分割数据增强——图像和标注同步增强

语义分割数据增强

常见的数据增强方式

查看pytorchtorchvision的transformer中的源代码，我们可以看到具有以下数据增强方式：

__all__ = ["Compose", "ToTensor", "PILToTensor", "ConvertImageDtype", "ToPILImage", "Normalize", "Resize", "Scale",
           "CenterCrop", "Pad", "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop",
           "RandomHorizontalFlip", "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop",
           "LinearTransformation", "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale",
           "RandomPerspective", "RandomErasing", "GaussianBlur", "InterpolationMode", "RandomInvert", "RandomPosterize",
           "RandomSolarize", "RandomAdjustSharpness", "RandomAutocontrast", "RandomEqualize"]

其中常见的数据增强方式包括：旋转、垂直翻转、水平翻转、放缩、剪裁、归一化等。

语义分割和图像分类的数据增强差异在于：语义分割是对图像的每个像素进行分类，所以在进行某些数据增强时，需要对标注图像(mask)进行同步操作，如旋转、剪裁、翻转等。

查遍网上的一些教程，但是没有发现一个能够直接使用的pytorch数据增强方式，所以想自己写一个方便后续使用。

具体实现代码

我们这里以一个细胞语义分割数据集为例，由于该数据集是灰度图像，所以相对于彩色图像数据增强有一些差距，代码中注释了灰度图像不能使用的数据增强方式，但是彩色图像可以使用的数据增强方式。具体代码如下所示：

import numpy as np
import cv2
import torch
from torch.utils.data import Dataset
import os
from PIL import Image
from torchvision.transforms import functional as F
import random


class CellDataset(Dataset):
    def __init__(self, image_dir, mask_dir, names_list, image_size=224, isGray=False, augmentation=True):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.augmentation = augmentation
        self.names_list = names_list
        self.isGray = isGray
        self.image_size = image_size

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

    def augmentate(self, image, mask):
        # it is expected to be in [..., H, W] format
        image = torch.unsqueeze(torch.from_numpy(
            np.array(image, dtype=np.uint8)), dim=0)
        mask = torch.unsqueeze(torch.from_numpy(
            np.array(mask, dtype=np.uint8)), dim=0)

        image = F.resize(image, size=[self.image_size, self.image_size])
        mask = F.resize(mask, size=[self.image_size, self.image_size])
        # 彩色图可以进行以下数据增强，参数不太好调整
        image = F.adjust_gamma(image, gamma=random.uniform(0.8, 1.2))
        image = F.adjust_contrast(
            image, contrast_factor=random.uniform(0.8, 1.2))
        image = F.adjust_brightness(
            image, brightness_factor=random.uniform(0.8, 1.2))
        image = F.adjust_saturation(
            image, saturation_factor=random.uniform(0.8, 1.2))
        image = F.adjust_hue(image, hue_factor=random.uniform(-0.2, 0.2))

        # 让image和mask进行同步旋转和翻转数据增强
        image_mask = torch.cat([image, mask], dim=0)

        if random.uniform(0, 1) > 0.5:
            image_mask = F.hflip(image_mask)
        if random.uniform(0, 1) > 0.5:
            image_mask = F.vflip(image_mask)
        if random.uniform(0, 1) > 0.5:
            image_mask = F.rotate(image_mask, angle=90)

        # 要看image和mask的维度
        image = image_mask[0, ...]
        mask = image_mask[1, ...]
        # image = image / 255
        # mask = mask / 255

        # image = torch.unsqueeze(image, dim=0)
        # 标准化，彩色图像需要传三个值
        # image = F.normalize(image, mean=[0.5], std=[1])
        # mask = torch.unsqueeze(mask, dim=0)

        return image, mask

    def __getitem__(self, item):
        image_path = os.path.join(self.image_dir, self.names_list[item])
        mask_path = os.path.join(self.mask_dir, self.names_list[item])
        image = Image.open(image_path)
        if self.isGray:
            image = image.convert('L')
        mask = Image.open(mask_path)

        if self.augmentation:
            image, mask = self.augmentate(image, mask)
        return image, mask


if __name__ == '__main__':
    cell_dataset = CellDataset(image_dir='./data/image', mask_dir='./data/label',
                               names_list=['0.png'])
    index = 3
    for image, mask in cell_dataset:
        print(image.shape, mask.shape)
        print(torch.max(image), torch.min(image))
        image = np.array(image, dtype=np.uint8)
        mask = np.array(mask, dtype=np.uint8)
        # cv2.imshow('image', image)
        # cv2.imshow('mask', mask)
        # cv2.waitKey(0)
        cv2.imwrite(os.path.join(
            './data/augment/image', str(index)+'.png'), image)
        cv2.imwrite(os.path.join(
            './data/augment/label', str(index)+'.png'), mask)

我们以下图图像和图像标注掩码为例进行实验：

原始细胞图像

原始细胞掩码标签图像

我们更该index的值，进行重复执行程序，来生成多个不同对应的数据增强图像，我们重复执行了4次，得到了以下数据增强的细胞图像和对应的掩码标签图像。

如果有问题可以在评论区进行回复。如果对您有帮助的话可以帮忙点赞👍👍👍。