高级实训任务2：U-Net图像分割

任务描述

将卷积神经网络（CNN）应用在图像分割任务上，我们需要对网络结构进行设计。
可供选择的网络有：YOLO v3、Mask R-CNN、U-Net。
在此项目中，我选择的是U-Net。
数据集选择的是 isbi_2012，即可选数据集网站的如下图项：

U-Net原理

什么是图像分割

图像分割就是要将一个图像分割成许多个部分，对于一个分割完毕的图像，每一个像素点都属于且只属于一个部分，同一个部分的所有像素点必须是相邻的，联通的。
因此不难看出，简单的图像检测只需要给出一个物体的大致位置就行，但是图像分割必须给出每一个部分的精确轮廓。
为此可以用一个函数去求出这个轮廓的位置和形状，可以用神经网络算法去求这个函数。

U-Net的先驱–FCN

FCN是第一个将全卷积网络与图像分割结合在一起的框架，FCN的框架模型如下图：

FCN采取解决方法是将pool4、pool3、和特征map融合起来，由于pool3、pool4、特征map大小尺寸是不一样的，所以融合应该前上采样到同一尺寸。这里的融合是拼接在一起，不是对应元素相加。

可以看出，由于FCN采取的是全卷积网络，在模型中对图像多次求卷积，在输出层得到的信息，也是对原图进行特征提取后的数据。但是图像分割想要得到的结果是一个和原图大小一样的图像，即一个描述各部分轮廓的黑白图。为了得到这个轮廓图像，必修将输出的特征信息还原成图像。
为此，FCN采取反卷积的方法来还原图像。将所得的特征上采样回去，再将预测结果做一一对应的分类，区分某一点在图像中的意义。因此在还原时，分割问题就转化成了分类问题。
这样会丢失很多信息。

U-Net

网络结构

U-Net的框架如下图：

U-Net分为两个部分，特征提取部分和上采样部分。

特征提取部分在图中包括左半边和下边。上采样部分包括右边。由于整个模型看起来像是U字型，所以称为U-Net。

收缩路径就是常规的卷积网络，它包含重复的2个3x3卷积，紧接着是一个RELU，一个max pooling（步长为2），用来降采样，每次降采样我们都将feature channel减半。扩展路径包含一个上采样（2x2上卷积），这样会减半feature channel，接着是一个对应的收缩路径的feature map，然后是2个3x3卷积，每个卷积后面跟一个RELU，因为每次卷积会丢失图像边缘，所以裁剪是有必要的，最后来一个1x1的卷积，用来将有64个元素的feature vector映射到一个类标签，整个网络一共有23个卷积层。

overlap-tile策略

overlap-tile策略是U-Net中使用到的一种优化策略。
由于医学影像很大，一般不能直接作为网络的输入，所以该策略会把训练集的图像分成若干小部分，对每一个部分进行训练。
在将大图像分块的时候，需要对每一个小块求卷积，但是图像边界的像素点没有周围像素，求卷积会导致信息的缺失，因此需要对划分好的每一个小块做一次扩充。采用镜像对称的原理，将矩形小块按其边界进行镜像对称，这样小块的四周会多处一圈和自己对称的部分，避免求卷积导致的信息丢失。

弹性变换策略

深度神经网络拥有很强的学习能力，因此如果训练数据集的内容不够，训练的数据集太小，会产生过拟合问题，导致训练好的模型无法正确使用。
因此在数据集有限的情况下，需要人为地对数据集进行扩充。比如将图像拉伸，裁剪，旋转，加入各种噪声等。
由于U-Net用于解决细胞组织图像的问题，根据实际情况，细胞组织的边界每时每刻都在做不规则的畸形变换。因此可以对图像加入这种噪声，在适量的范围内，随即改变细胞的边界，改变训练集的图像，从而扩展训练集数据。

实现方法

代码见：这里

工具选择

语言：python
python语言方便快捷，出错率小，适合做复杂的算法。

框架选择：tensorflow，keras

其中tensorflow是主要框架，keras只引用其中的图像处理功能，辅助tensorflow完成任务。

此外还用了absl库函数。用absl库定义可重复的代码段，定义训练的各个参数（比如学习率，周期数，每个周期的训练轮数等）。

参数设定：eopch=3（训练三个周期）
step=100（每个周期训练100轮）
batch_size=2（每次训练取两个样本）
learning_rate=0.0001（学习率）

tensorflow定义网络

如下是代码中关于网络框架的定义。

inputs = tf.keras.layers.Input((512, 512, 1))

        # Contracting part
        conv1 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
        conv1 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
        assert conv1.shape[1:] == (512, 512, 64)
        pool1 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv1)
        assert pool1.shape[1:] == (256, 256, 64)
        conv2 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
        conv2 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
        assert conv2.shape[1:] == (256, 256, 128)
        pool2 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv2)
        assert pool2.shape[1:] == (128, 128, 128)
        conv3 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
        conv3 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
        assert conv3.shape[1:] == (128, 128, 256)
        pool3 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv3)
        assert pool3.shape[1:] == (64, 64, 256)
        conv4 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
        conv4 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
        drop4 = tf.keras.layers.Dropout(0.5)(conv4)
        assert drop4.shape[1:] == (64, 64, 512)
        pool4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(drop4)
        assert pool4.shape[1:] == (32, 32, 512)

        conv5 = tf.keras.layers.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(
            pool4)
        conv5 = tf.keras.layers.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(
            conv5)
        assert conv5.shape[1:] == (32, 32, 1024)
        drop5 = tf.keras.layers.Dropout(0.5)(conv5)

        # Expansive part
        up6 = tf.keras.layers.Conv2D(512, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
            tf.keras.layers.UpSampling2D(size=(2, 2))(drop5))
        assert up6.shape[1:] == (64, 64, 512)
        merge6 = tf.keras.layers.concatenate([drop4, up6], axis=3)
        assert merge6.shape[1:] == (64, 64, 1024)
        conv6 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(
            merge6)
        conv6 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
        assert conv6.shape[1:] == (64, 64, 512)

        up7 = tf.keras.layers.Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
            tf.keras.layers.UpSampling2D(size=(2, 2))(conv6))
        assert up7.shape[1:] == (128, 128, 256)
        merge7 = tf.keras.layers.concatenate([conv3, up7], axis=3)
        assert merge7.shape[1:] == (128, 128, 512)
        conv7 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(
            merge7)
        conv7 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)
        assert conv7.shape[1:] == (128, 128, 256)

        up8 = tf.keras.layers.Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
            tf.keras.layers.UpSampling2D(size=(2, 2))(conv7))
        assert up8.shape[1:] == (256, 256, 128)
        merge8 = tf.keras.layers.concatenate([conv2, up8], axis=3)
        assert merge8.shape[1:] == (256, 256, 256)
        conv8 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(
            merge8)
        conv8 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)
        assert conv8.shape[1:] == (256, 256, 128)

        up9 = tf.keras.layers.Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
            tf.keras.layers.UpSampling2D(size=(2, 2))(conv8))
        assert up9.shape[1:] == (512, 512, 64)
        merge9 = tf.keras.layers.concatenate([conv1, up9], axis=3)
        assert merge9.shape[1:] == (512, 512, 128)
        conv9 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
        conv9 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
        assert conv9.shape[1:] == (512, 512, 64)
        conv9 = tf.keras.layers.Conv2D(2, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
        assert conv9.shape[1:] == (512, 512, 2)
        conv10 = tf.keras.layers.Conv2D(num_classes, 1, activation='sigmoid')(conv9)
        assert conv10.shape[1:] == (512, 512, num_classes)

        model = tf.keras.Model(inputs=inputs, outputs=conv10)

其中conv1~10是上文图中的卷积层的卷积的定义。pool序列是池化层的定义。

训练

我使用的isbi_2012数据集有如下的结构：

分别为测试集，训练集，训练集的标签。

用如下代码开始训练。

def make_train_generator(batch_size, aug_dict):
    image_gen = ImageDataGenerator(**aug_dict)
    mask_gen = ImageDataGenerator(**aug_dict)

    # set image and mask same augmentation using same seed
    image_generator = image_gen.flow_from_directory(
        directory='./isbi_2012/preprocessed',
        classes=['train_imgs'],
        class_mode=None,
        target_size=(512, 512),
        batch_size=batch_size,
        color_mode='grayscale',
        seed=1
    )
    mask_generator = mask_gen.flow_from_directory(
        directory='./isbi_2012/preprocessed',
        classes=['train_labels'],
        class_mode=None,
        target_size=(512, 512),
        batch_size=batch_size,
        color_mode='grayscale',
        seed=1
    )
    train_generator = zip(image_generator, mask_generator)
    for (batch_images, batch_labels) in train_generator:
        batch_images, batch_labels = normolize(batch_images, batch_labels)

        yield (batch_images, batch_labels)

这段代码给出了训练集和训练集标签的读取，定义了目标文件的大小（512*512）和读取方式（灰度图），由于算法中需要采用随机数，两者都用同样的随机种子（seed=1）。

结果

经过了1小时的训练，运行结果如下图：

随着周期数的推进，训练准确度也越来越高。但是我的计算资源有限，论文中训练了10个周期，每个周期2000轮，准确率会非常高。