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

1. 介绍

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

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

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

文件目录如下:

Unet 实战分割项目、多尺度训练、多类别分割_第1张图片

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灰度值的方法,如下:

Unet 实战分割项目、多尺度训练、多类别分割_第2张图片

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

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

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

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

Unet 实战分割项目、多尺度训练、多类别分割_第3张图片

保存的txt文本如下:

这是二值图像分割

Unet 实战分割项目、多尺度训练、多类别分割_第4张图片

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

Unet 实战分割项目、多尺度训练、多类别分割_第5张图片

2.2 加载mask 灰度值映射

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

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

Unet 实战分割项目、多尺度训练、多类别分割_第6张图片

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

Unet 实战分割项目、多尺度训练、多类别分割_第7张图片

这部分代码就是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 加载数据可视化

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

Unet 实战分割项目、多尺度训练、多类别分割_第8张图片

如下:

Unet 实战分割项目、多尺度训练、多类别分割_第9张图片

对应的灰度值:

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

2.4 预测脚本

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

3. 代码使用

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

Unet 实战分割项目、多尺度训练、多类别分割_第10张图片

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

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

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

Unet 实战分割项目、多尺度训练、多类别分割_第11张图片

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

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

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

Unet 实战分割项目、多尺度训练、多类别分割_第12张图片

3.1 训练过程

最后两个epoch结果

Unet 实战分割项目、多尺度训练、多类别分割_第13张图片

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

Unet 实战分割项目、多尺度训练、多类别分割_第14张图片

loss 曲线:

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

Unet 实战分割项目、多尺度训练、多类别分割_第15张图片

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

Unet 实战分割项目、多尺度训练、多类别分割_第16张图片

Unet 实战分割项目、多尺度训练、多类别分割_第17张图片

Unet 实战分割项目、多尺度训练、多类别分割_第18张图片

3.2注意的点

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

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

3.3 代码下载

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

4. 未来展望

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

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

你可能感兴趣的:(图像分割,深度学习,人工智能,机器学习)