Unet 实战分割项目、多尺度训练、多类别分割

1. 介绍

之前写了篇二值图像分割的项目，支持多尺度训练，网络采用backbone为vgg的unet网络。缺点就是没法实现多类别的分割，具体可以参考：二值图像分割统一项目

本章只对增加的代码进行介绍，其余的参考上述链接博文

本章实现的unet网络的多类别分割，也就是分割可以是两个类别，也可以是多个类别。训练过程仍然采用多尺度训练，即网络会随机将图片缩放到设定尺寸的0.5-1.5倍之间

文件目录如下：

2. 实现思路

因为多类别的分割，mask模板都是灰度图，一般0为背景，255为前景。多类别的话，就是0为背景，1 ， 2，3等等其他的灰度值为前景。

而为了方便观察，前景的灰度值不会设定的这么近，都会相隔很远。例如4分割，不会前景分为0123，而是0，63，127，255类似这样的

值得注意的是，VOC数据集的mask模板也是灰度图像，只是填充的颜色才导致显示出来是彩色的

2.1 mask的灰度值

多分割的unet网络，输出的通道数就是分割的个数，从0开始依次递增，所以将mask的灰度映射成0 1 2 3 是必要的

这个就类似于二值图像，将前景 255 映射为 1

因为对于不同的任务，mask的灰度值是不同的，对于新人小白也不会查看mask的灰度值或者找不全等等。要是自己在dataset里面一个个映射也很麻烦，这里提供了一个自动实现的方法

在utils脚本里提供了一个查找mask灰度值的方法，如下：

np.unique 是查找数组里出现过的数字，例如0 0 1 2 2，返回的就是0 1 2 

这里将gray 内容按照从小到大排序，后面映射需要用！！！

代码会遍历所有的训练图像的mask，找到所有的mask前景+背景（0）灰度值，为了在dataset和预测脚本使用，这里将前景的分割像素点灰度值保存为txt格式，保存的路径是data文件夹下

这里返回的len(gray)就是分割的个数，包括背景的，这样接受compute_gray 函数的返回值，可以直接定义unet 分割的通道数

保存的txt文本如下：

这是二值图像分割

这是腹部多脏器多类别分割

2.2 加载mask 灰度值映射

这部分内容在dataset.py脚本中

首先加载txt文本，self.txt 是紫色框中的列表内容

这里很常见，通过image找到mask图片，用replace根据自己图片后缀替换即可

这部分代码就是mask的灰度值映射

首先将当前mask的所有前景找到，用gray遍历。因为之前txt的灰度值是从小到大且从背景0开始排序的，而且self.txt加载的txt是列表形式。可以取个巧，将index就作为分割的映射值

例如，txt 内容是 0 62 125 252 ，说明这是一个包含背景四分类的分割项目

那么self.txt 列表的内容就是【0 62 125 252】

我们想要映射的结果就是0-->0、62-->1、125-->2、252-->3，而0123这些不正是列表值的index吗

2.3 加载数据可视化

打开这部分代码，可以查看可视化数据

如下：

对应的灰度值：

注意：这里classes打印出来有255，不是映射失败，因为本章采用的多尺度训练，图像放大后，会用255填充，这是没关心的，因为后面计算交叉熵损失会忽略255的像素

2.4 预测脚本

预测也很简单，加载完，在映射回去即可

3. 代码使用

说了这么多，代码如何使用最重要，README如下

这里需要更改的地方就三个：

1. 自定义数据按照data目录摆放，看README第二点

2. 因为image个mask的图像后缀不一定严格一致，所以要根据自己的数据更改dataset代码

3. 为了多尺度训练效果可以更好，建议将下面两个参数尽量改成和数据size接近的

base-size 为多尺度的尺寸，图像会缩放成 0.5 * base-size --- 1.5 * base-size之间

crop-size 会中心裁剪成规定的，建议改成和数据宽高接近的

3.1 训练过程

最后两个epoch结果

训练日志有每个类别的recall和precision：

loss 曲线：

网络还没有收敛，只是测试，增大epoch可以收敛的更好

推理：依次为原图、推理图、真实GT图片

3.2注意的点

需要注意的是，有的数据mask不是多值图像，例如二分割任务，像素点除了0 255仍有中间的灰度值等等。这时候建议检查数据，通过opencv处理在进行训练，否则会出错

除此外，因为这里的255也会被映射回去，作为分割的一部分。所以类似于VOC这种有255忽略点的数据集可能不支持本项目，当然可以自己更改代码，将dataset中255不映射即可，或者把txt文本中255删除

3.3 代码下载

项目封装在这：深度学习 Unet 实战分割项目、多尺度训练、多类别分割：腹部多脏器5类别分割数据集

4. 未来展望

后续，会将resnet加入unet中，对比下效果。还有更多的高阶API分割模型等等，也会做个项目

这里展示的是腹部多脏器MRI的多分割项目，代码对二值图像DRIVE数据集也做了测试，不必担心不能兼容