U-Net 的解读与编码实现（Tensorflow 2）

汉北集团 · 发表于 2023-1-8 13:44:24

1、文章相关

题目：U-Net: Convolutional Networks for Biomedical Image Segmentation
作者：Olaf Ronneberger, Philipp Fischer, and Thomas Brox
发表时间：2015年
主要贡献：医疗图像分割的利器与开山之作、对FCN进行改进、转置卷积
U-Net网络非常简单，前半部分作用是特征提取，后半部分是上采样，通常将这样的结构叫做编码器-解码器（Encoder-Decoder）结构。由于此网络整体结构类似于大写的英文字母U，故得名U-Net。
U-Net与FCN网络有一点非常不同的地方：

U-Net在上采样使用的是转置卷积（也有人称反卷积），而FCN上采样使用的是向上池化操作。
U-Net在实现Skip Connection时采用拼接的方式进行特征融合（通道数叠加），而FCN融合时使用的对应点相加（通道数不变）。

2、网络架构分析

U-Net网络架构

输入模块I（64@568×568）：

输入（3@572×572）：输入图像大小为572×572，三通道。
卷积层I_C_1（64@570×570）：使用64通道大小为3×3的卷积核对输入图像卷积计算得到64个大小为570×570的特征图。
ReLu：使用ReLU激活函数。
卷积层I_C1_2（64@568×568）：使用64通道大小为3×3的卷积核对输入图像卷积计算得到64个大小为568×568的特征图。
ReLu：使用ReLU激活函数。

收缩模块C1（128@280×280）：

MaxPooling（64@284×284）：使用最大池化，池化单元的规格为2×2，步长为2，池化后的结果为64@284×284的特征图。
卷积层C_C1_1（128@282×282）：使用128通道大小为3×3的卷积核对输入图像卷积计算得到128个大小为282×282的特征图。
ReLu：使用ReLU激活函数。
卷积层C_C1_2（128@280×280）：使用128通道大小为3×3的卷积核对输入图像卷积计算得到128个大小为280×280的特征图。
ReLu：使用ReLU激活函数。

收缩模块C2（256@136×136）：

MaxPooling（128@140×140）：使用最大池化，池化单元的规格为2×2，步长为2，池化后的结果为128@140×140的特征图。
卷积层C_C2_1（256@138×138）：使用256通道大小为3×3的卷积核对输入图像卷积计算得到256个大小为138×138的特征图。
ReLu：使用ReLU激活函数。
卷积层C_C2_2（256@136×136）：使用256通道大小为3×3的卷积核对输入图像卷积计算得到256个大小为136×136的特征图。
ReLu：使用ReLU激活函数。

收缩模块C3（512@64×64）：

MaxPooling（256@68×68）：使用最大池化，池化单元的规格为2×2，步长为2，池化后的结果为256@68×68的特征图。
卷积层C_C3_1（512@66×66）：使用512通道大小为3×3的卷积核对输入图像卷积计算得到512个大小为66×66的特征图。
ReLu：使用ReLU激活函数。
卷积层C_C3_2（512@64×64）：使用512通道大小为3×3的卷积核对输入图像卷积计算得到512个大小为64×64的特征图。
ReLu：使用ReLU激活函数。

收缩模块C4（1024@28×28）

MaxPooling（512@32×32）：使用最大池化，池化单元的规格为2×2，步长为2，池化后的结果为512@32×32的特征图。
卷积层C_C4_1（1024@30×30）：使用1024通道大小为3×3的卷积核对输入图像卷积计算得到1024个大小为30×30的特征图。
ReLu：使用ReLU激活函数。
卷积层C_C4_2（1024@28×28）：使用1024通道大小为3×3的卷积核对输入图像卷积计算得到1024个大小为28×28的特征图。
ReLu：使用ReLU激活函数。

扩张模块E4（512@52×52）

上采样层E_TC_4（512@56×56）：上采样使用转置卷积的方式，使用512个卷积核，转置卷积的卷积核大小为2×2，步长为2，得到512个大小为56×56的特征图。
特征图拼接（Skip Connection）（1024@56×56）：先将收缩模块C3的特征结果（大小为512@64×64）裁剪为512@56×56，再与上一步中上采样结果（512@56×56）进行拼接（通道叠加），得到1024@56×56。
卷积层E_C4_1（512@54×54）：使用512通道大小为3×3的卷积核对输入图像卷积计算得到512个大小为54×54的特征图。
ReLu：使用ReLU激活函数。
卷积层E_C4_2（512@52×52）：使用512通道大小为3×3的卷积核对输入图像卷积计算得到512个大小为52×52的特征图。
ReLu：使用ReLU激活函数。

扩张模块E3（256@100×100）

上采样层E_TC_3（256@104×104）：上采样使用转置卷积的方式，使用256个卷积核，转置卷积的卷积核大小为2×2，步长为2，得到256个大小为104×104的特征图。
特征图拼接（Skip Connection）（512@104×104）：先将收缩模块C2的特征结果（大小为256@136×136）裁剪为256@104×104，再与上一步中上采样结果（256@104×104）进行拼接（通道叠加），得到512@104×104。
卷积层E_C3_1（256@102×102）：使用256通道大小为3×3的卷积核对输入图像卷积计算得到256个大小为102×102的特征图。
ReLu：使用ReLU激活函数。
卷积层E_C3_2（256@100×100）：使用256通道大小为3×3的卷积核对输入图像卷积计算得到256个大小为100×100的特征图。
ReLu：使用ReLU激活函数。

扩张模块E2（128@196×196）

上采样层E_TC_2（128@200×200）：上采样使用转置卷积的方式，使用128个卷积核，转置卷积的卷积核大小为2×2，步长为2，得到128个大小为200×200的特征图。
特征图拼接（Skip Connection）（128@200×200）：先将收缩模块C1的特征结果（大小为128@280×280）裁剪为128@200×200，再与上一步中上采样结果（128@200×200）进行拼接（通道叠加），得到256@200×200。
卷积层E_C2_1（128@198×198）：使用128通道大小为3×3的卷积核对输入图像卷积计算得到128个大小为198×198的特征图。
ReLu：使用ReLU激活函数。
卷积层E_C2_2（128@196×196）：使用128通道大小为3×3的卷积核对输入图像卷积计算得到128个大小为196×196的特征图。
ReLu：使用ReLU激活函数。

扩张模块E1（64@388×388）

上采样层E_TC_2（128@200×200）：上采样使用转置卷积的方式，使用128个卷积核，转置卷积的卷积核大小为2×2，步长为2，得到128个大小为200×200的特征图。
特征图拼接（Skip Connection）（128@392×392）：先将输入模块I的特征结果（大小为64@568×568）裁剪为64@392×392，再与上一步中上采样结果（64@392×392）进行拼接（通道叠加），得到128@392×392。
卷积层E_C2_1（64@390×390）：使用64通道大小为3×3的卷积核对输入图像卷积计算得到64个大小为390×390的特征图。
ReLu：使用ReLU激活函数。
卷积层E_C2_2（64@388×388）：使用64通道大小为3×3的卷积核对输入图像卷积计算得到64个大小为388×388的特征图。
ReLu：使用ReLU激活函数。

输出模块O（2@288×388）：使用2通道大小为1×1的卷积核对输入图像卷积计算得到2个大小为388×388的特征图。
3、代码实现

数据集：data-science-bowl-2018
框架：TF2
复现代码是使用的数据集是data-science-bowl-2018，图片大小为256×256×3，训练集样本数目为670个
网络搭建的具体思路为：四个下采样模块后，进行四个上采样模块，上采样是需要进行特征通道的拼接，分解后的主要操作如下=>输入->卷积->relu->卷积->relu->最大池化->卷积->relu->卷积->relu->最大池化->卷积->relu->卷积->relu->最大池化->卷积->relu->卷积->relu->最大池化->卷积->relu->卷积->relu->上采样->拼接特征->卷积->relu->卷积->relu->上采样->拼接特征->卷积->relu->卷积->relu->上采样->拼接特征->卷积->relu->卷积->relu->上采样->拼接特征->卷积->relu->卷积->relu->卷积->输出。
引入模块
import tensorflow as tf
from tensorflow import keras
import os
import numpy as np
from tqdm import tqdm
from skimage.io import imread, imshow
from skimage.transform import resize
import matplotlib.pyplot as plt
import random参数定义
seed = 42
np.random.seed = seed

# 样本图片大小
IMG_WIDTH = 256
IMG_HEIGHT = 256
IMG_CHANNELS = 3

# 数据集路径
DATA_PATH = &#39;data/bowl2018/&#39;注意调整数据集的路径，默认在根目录下
数据加载和处理
# 数据加载
TRAIN_PATH = DATA_PATH + &#39;stage1_train/&#39;  # 训练集路径
TEST_PATH = DATA_PATH + &#39;stage1_test/&#39;  # 测试集路径

train_ids = next(os.walk(TRAIN_PATH))[1]
test_ids = next(os.walk(TEST_PATH))[1]

# 构造训练集输入和输出（mask）
X_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
Y_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
path = TRAIN_PATH + id_
img = imread(path + &#39;/images/&#39; + id_ + &#39;.png&#39;)[:, :, :IMG_CHANNELS]
img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode=&#39;constant&#39;, preserve_range=True)
X_train[n] = img  # Fill empty X_train with values from img
mask = np.zeros((IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
# mask
for mask_file in next(os.walk(path + &#39;/masks/&#39;))[2]:  # os.walk()文件、目录遍历器,在目录树中游走输出在目录中的文件名
      mask_ = imread(path + &#39;/masks/&#39; + mask_file)
      mask_ = np.expand_dims(resize(mask_, (IMG_HEIGHT, IMG_WIDTH), mode=&#39;constant&#39;, preserve_range=True), axis=-1)
      mask = np.maximum(mask, mask_)

Y_train[n] = mask

# 构造测试集输入
X_test = np.zeros((len(test_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
for n, id_ in tqdm(enumerate(test_ids), total=len(test_ids)):
path = TEST_PATH + id_
img = imread(path + &#39;/images/&#39; + id_ + &#39;.png&#39;)[:, :, :IMG_CHANNELS]
img = resize(img, (IMG_HEIGHT, IMG_WIDTH), mode=&#39;constant&#39;, preserve_range=True)
X_test[n] = img

数据处理进度条

样本图片预览
# 画出前25张图
plt.figure(figsize=(20,20))
for i in range(25):
plt.subplot(5,5,i+1)  # 每行五列，共五行
plt.xticks([])
plt.yticks([])
plt.grid(False)
plt.imshow(X_test, cmap=plt.cm.binary)
plt.show()

前25个样本的可视化

网络模块实现以及结构搭建
inputs = tf.keras.layers.Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
s = tf.keras.layers.Lambda(lambda x: x / 255)(inputs)输入模块（第一个模块）
# UNet输入模块
def InputBlock(input, filters, kernel_size=3, strides=1, padding=&#39;same&#39;):
conv_1 = tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                                 activation=&#39;relu&#39;)(input)  # 卷积块1
return tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                              activation=&#39;relu&#39;)(conv_1)  # 卷积块2下采样模块
# 收缩路径模块
def ContractingPathBlock(input, filters, kernel_size=3, strides=1, padding=&#39;same&#39;):
down_sampling = tf.keras.layers.MaxPool2D((2, 2))(input)  # 最大池化
conv_1 = tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                                 activation=&#39;relu&#39;)(down_sampling)  # 卷积块1
return tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                                 activation=&#39;relu&#39;)(conv_1)  # 卷积块2上采样模块
# 扩张（恢复）路径模块
def ExpansivePathBlock(input, con_feature, filters, tran_filters, kernel_size=3, tran_kernel_size=2, strides=1,
                     tran_strides=2, padding=&#39;same&#39;, tran_padding=&#39;same&#39;):
upsampling = tf.keras.layers.Conv2DTranspose(filters=tran_filters, kernel_size=tran_kernel_size,
                                             strides=tran_strides, padding=tran_padding)(input)  # 上采样（转置卷积方式）
con_feature = tf.image.resize(con_feature, ((upsampling.shape)[1], (upsampling.shape)[2]),
                              method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)  # 裁剪需要拼接的特征图
concat_feature = tf.concat([con_feature, upsampling], axis=3)  # 拼接扩张层和收缩层的特征图（skip connection）
conv_1 = tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                                 activation=&#39;relu&#39;)(concat_feature)  # 卷积1
return tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding,
                              activation=&#39;relu&#39;)(conv_1)  # 卷积2U-Net网络架构构建
# UNet网络架构
def UNet(input_shape):
inputs = tf.keras.layers.Input(input_shape)
s = tf.keras.layers.Lambda(lambda x: x / 255)(inputs)

# input block
input_block = InputBlock(s, 64)

# contracting path
con_1 = ContractingPathBlock(input_block, 128)
con_2 = ContractingPathBlock(con_1, 256)
con_3 = ContractingPathBlock(con_2, 512)
con_4 = ContractingPathBlock(con_3, 1024)

# expansive path
exp_4 = ExpansivePathBlock(con_4, con_3, 512, 512)
exp_3 = ExpansivePathBlock(exp_4, con_2, 256, 256)
exp_2 = ExpansivePathBlock(exp_3, con_1, 128, 128)
exp_1 = ExpansivePathBlock(exp_2, input_block, 64, 64)

outputs = tf.keras.layers.Conv2D(2, 1)(exp_1)  # 最终输出

return tf.keras.Model(inputs=[inputs], outputs=[outputs])训练
model = UNet(input_shape=(IMG_WIDTH, IMG_HEIGHT, IMG_CHANNELS))
model.summary()
model.compile(optimizer=&#39;adam&#39;, loss=&#39;binary_crossentropy&#39;, metrics=[&#39;accuracy&#39;])
results = model.fit(X_train, Y_train, batch_size=10, epochs=10)

经过10个epoch之后的结果

损失
# 损失下降曲线
plt.plot(results.history[&#39;loss&#39;])
plt.title(&#39;model loss&#39;)
plt.ylabel(&#39;loss&#39;)
plt.xlabel(&#39;epoch&#39;)
plt.show()

10个epoch后的loss下降曲线

4、补充说明

精简版的代码已上传github，可直接在pycharm中运行，自取地址：

如果运行时提示内存不够，可以将图片尺寸调整为128×128，减小batch_size

骚 · 发表于 2023-1-8 13:44:42

您好，使用next函数读取的时候会出现StopIteration，这个问题该怎么解决呢，感谢您

郭联 · 发表于 2023-1-8 13:44:58

你好，数据集可以分享下吗？

夵尖 · 发表于 2023-1-8 13:45:24

https://www.kaggle.com/c/data-science-bowl-2018/data
我是从这里下载的，保存之后更改一下代码里的路径就行

晓珂侃世界 · 发表于 2023-1-8 13:45:58

最近领插值法那里，应该是把64x64直接resize成56x56,最后的结果才会是1x388x388x1.

懒得娶姳 · 发表于 2023-1-8 13:46:40

您好，如果是自己的数据集怎么导入呢，我是直接把train分为images和masks两个文件夹的

新依诺广告 · 发表于 2023-1-8 13:46:47

MaxPooling为什么会修改通道数呢？是不是搞错了

无敌霸气你冰哥 · 发表于 2023-1-8 13:47:22

你好，请问这个问题你解决了吗，我也遇到这个情况了

暖心 · 发表于 2023-1-8 13:47:46

谢谢你的提醒，这里确实出现了错误，已经修正了

樊斌 · 发表于 2023-1-8 13:48:43

最后一层,不是应该是2通道, 卷积层大小1x1嘛,
那outputs = tf.keras.layers.Conv2D(1, 1)(exp_1) # 最终输出
应该为 outputs = tf.keras.layers.Conv2D(2, 1)(exp_1) # 最终输出不然通道数不就是1了吗.... 是我弄错了吗

U-Net 的解读与编码实现（Tensorflow 2）

浏览过的版块