之前写了篇二值图像分割的项目,支持多尺度训练,网络采用backbone为vgg的unet网络。缺点就是没法实现多类别的分割,具体可以参考:二值图像分割统一项目
本章只对增加的代码进行介绍,其余的参考上述链接博文
本章实现的unet网络的多类别分割,也就是分割可以是两个类别,也可以是多个类别。训练过程仍然采用多尺度训练,即网络会随机将图片缩放到设定尺寸的0.5-1.5倍之间
文件目录如下:
因为多类别的分割,mask模板都是灰度图,一般0为背景,255为前景。多类别的话,就是0为背景,1 , 2,3等等其他的灰度值为前景。
而为了方便观察,前景的灰度值不会设定的这么近,都会相隔很远。例如4分割,不会前景分为0123,而是0,63,127,255类似这样的
值得注意的是,VOC数据集的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文本如下:
这是二值图像分割
这是腹部多脏器多类别分割
这部分内容在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吗
打开这部分代码,可以查看可视化数据
如下:
对应的灰度值:
注意:这里classes打印出来有255,不是映射失败,因为本章采用的多尺度训练,图像放大后,会用255填充,这是没关心的,因为后面计算交叉熵损失会忽略255的像素
预测也很简单,加载完,在映射回去即可
说了这么多,代码如何使用最重要,README如下
这里需要更改的地方就三个:
1. 自定义数据按照data目录摆放,看README第二点
2. 因为image个mask的图像后缀不一定严格一致,所以要根据自己的数据更改dataset代码
3. 为了多尺度训练效果可以更好,建议将下面两个参数尽量改成和数据size接近的
base-size 为多尺度的尺寸,图像会缩放成 0.5 * base-size --- 1.5 * base-size之间
crop-size 会中心裁剪成规定的,建议改成和数据宽高接近的
最后两个epoch结果
训练日志有每个类别的recall和precision:
loss 曲线:
网络还没有收敛,只是测试,增大epoch可以收敛的更好
推理:依次为原图、推理图、真实GT图片
需要注意的是,有的数据mask不是多值图像,例如二分割任务,像素点除了0 255仍有中间的灰度值等等。这时候建议检查数据,通过opencv处理在进行训练,否则会出错
除此外,因为这里的255也会被映射回去,作为分割的一部分。所以类似于VOC这种有255忽略点的数据集可能不支持本项目,当然可以自己更改代码,将dataset中255不映射即可,或者把txt文本中255删除
项目封装在这:深度学习 Unet 实战分割项目、多尺度训练、多类别分割:腹部多脏器5类别分割数据集
后续,会将resnet加入unet中,对比下效果。还有更多的高阶API分割模型等等,也会做个项目
这里展示的是腹部多脏器MRI的多分割项目,代码对二值图像DRIVE数据集也做了测试,不必担心不能兼容