pytorch结合albumentations在图像分割中的应用

前言

  在训练网络模型的时候经常会数据集不够导致网络模型出现过拟合，其中一个重要的解决方法就是增加数据集，但是当我们的数据集很难扩充，或者制作成本比较高的时候怎么办呢？这个时候就需要用到数据增强了，通过变换我们原始的数据集(如，旋转，对称，加噪声，改变对比度，改变亮度等)。我们知道pytorch框架自带的有transforms，通过transforms可以对原始数据进行一些变换，但是transforms自带的变换是有限的且transforms.compose里面进行的变换是从上到下顺序执行的，也就是说我们使用pytorch自带的数据增强只能选择一种增强方式，就很不好用。这个时候就到了albumentations出场的时候了(也有其他的库，如imgaug,这里我们只介绍albumentations)，albumentations里面的OneOf可以让我们根据自己的需要选择某种数据增强出现的概率，从而达到扩充我们数据集的作用。
  还有就是我们在做分割或者检测的时候不仅仅要对image进行数据增强，对其相应的label也要做相同的变换，但是这个操作在torch自自带的transforms里面很不好用，往往需要自己来写函数实现就很麻烦。而albumentations只需要提供一个接口就能轻松完成。

一. 安装

很简单,直接pip就行了，通过清华镜像安装会比较快。

pip install albumentations -i https://pypi.tuna.tsinghua.edu.cn/simple

二.使用方法

  下面介绍一些albumentations常用的变换，具体课参考这个博主写的博文，里面列的很详细，我在此就不做介绍了：albumentations常用图像增强方法，也可以自己去albumentations的官网去搜索相关函数的用法。albumentations官网。
  下面我们以图像分割为例来简单使用下albumentations的用法。首先我们准备两张图像，一张image，一张label，如下图：


完整代码如下：

import numpy as np
import cv2
from PIL import Image
import albumentations as A
from matplotlib import pyplot as plt

image_path = "./data/Image/ILSVRC2012_test_00000004.jpg"
gt_path = "./data/GT/ILSVRC2012_test_00000004.png"
image = Image.open(image_path)
gt = Image.open(gt_path)

transform = A.Compose([
    A.Resize(256, 256, ),  # 裁剪
    A.RandomCrop(120, 120),
    A.Rotate(limit=40, p=0.9),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.3),
    A.OneOf([
        A.Blur(blur_limit=3, p=0.3),
        A.ColorJitter(p=0.5),
    ], p=1.0)
])

image_list = []
image = np.array(image)
mask = np.array(gt)
for i in range(10):
    augmentations = img_trans(mask=mask, image=image)
    augmented_img = augmentations["image"]
    augmented_gt = augmentations["mask"]
    image_list.append(augmented_img)
    image_list.append(augmented_gt)
for i in range(20):
    plt.imshow(image_list[i])
    plt.show()

重点解释下compose和OneOf的区别：

• Compose：自上而下顺序执行
• OneOf：随机选择一种执行

语句后面的p=0.5，表示每次进行此变换的概率。
重点注意事项：
下面是容易犯错的重点：mask，mask，mask，

augmentations = img_trans(mask=mask, image=image)

  此语句里面的mask=mask，前面一个mask是关键字，千万不能修改，不然image进行了数据增强，label不会发生任何变化。下面展示一些增强结果。






以上就是对albumentations的一个介绍，如有错误，敬请指正！