用Unet实现图像分割（by pytorch）

Segmentation

Figure1来自CamVid database，专为目标识别（Object Dection）和图像分割（Image Segmentation）提供训练数据的网站。从图中可以看出，segmentation将图像中不同个体用不同颜色来标记，这里不同的颜色就代表不同的分类，例如红色就是分类1，蓝色就是分类2，可以看出，它就是像素级的图像识别（Image Identification）。

除了自动驾驶之外，图像分割还广泛应用于医学诊断、卫星影像定位、图片合成等领域，本文就以当前kaggle上最热门的segmentation竞赛–TGS Salt Identification Challenge为例来讲解如何应用Unet来解决真实世界的图像分割问题。github: here。

TGS公司通过地震波反射技术绘制出下图所示的3D地质图像，并标记出图像中的盐矿区域，参赛者需要训练用于从岩层中分离盐矿的机器学习模型。

Figure 2是trainset中5组image和mark图片，每组的左图为原始地质图像，右图为该图像的分类，称为mark，黑色区域代表一般岩层，白色区域就是盐的分布。segmentation要做的就是训练一个image-to-image的模型，通过对原始图像的学习，生成其对应的mask₂，mask则作为target，通过最小化mask和mask₂的差距来识别哪些是盐。

Dataset

生成dataset的第一步是根据run length数据创建对应的mark图片，因为TGS的trainset里面已经提供了mark图片（mark图片和对应image图片同名），所以我们就不需要额外再创建。

但要知道的是，并非所有的segmentation dataset都会提供marks，你需要根据数据run length来为images创建相应的marks，run length是如下图rle_mask所示的数据，数据间以空格分隔，两两为一组，每组的第一个数代表flatten后的image vector的起始下标，后一个数代表它所占据的长度，占据区域会填充该目标对应的分类号，如0、1、2…，通过rle_decode()可以将run length转化为mark。

def rle_decode(mask_rle, shape=(101, 101)):
    s = mask_rle.split()
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.mean()

从Figure 2可以看到，地质图像都是低分辨画质，只有101x101大小，不仅不利于神经网络的卷积计算，也不利于图像识别，所以我们接下来一般会将其resize为128x128大小。

def resize_img(fn, outp_dir, sz):
  Image.open(fn).resize((sz, sz)).save(outp_dir/fn.name)

Data augmentation是创建dataset的核心，和object dection一样，segmentation一般不会做random crop，我在这个项目中采用水平、垂直翻转和微调光暗的方法来做augmentation。

aug_tfms = [
    RandomFlip(tfm_y=TfmType.CLASS),
    RandomDihedral(tfm_y=TfmType.CLASS),
#