Unet 多分类分割，附开源代码

Unet图像分割在大多的开源项目中都是针对于二分类，理论来说，对于多分类问题，依旧可行。可小编尝试过很多的方法在原有的开源代码进行多分类，分割的效果都不尽如意。最后下定决心自行修改网络，修改后的代码开源在github上（文末有链接），欢迎star。最终分割效果根据自己项目达到了要求，所以有需要的小伙伴可以尝试我们的方法。

Unet介绍

U-net前半部分作用是特征提取，后半部分是上采样。在一些文献中也把这样的结构叫做编码器-解码器结构。由于此网络整体结构类似于大写的英文字母U，故得名U-net。
U-net与其他常见的分割网络有一点非常不同的地方：U-net采用了完全不同的特征融合方式：拼接，U-net采用将特征在channel维度拼接在一起，形成更厚的特征。而FCN融合时使用的对应点相加，并不形成更厚的特征。

所以语义分割网络在特征融合时有两种办法：

1. FCN式的对应点相加，对应于TensorFlow中的tf.add()函数；
2. U-net式的channel维度拼接融合，对应于TensorFlow的tf.concat()函数，比较占显存。

除了上述新颖的特征融合方式，U-net还有以下几个优点：
1、5个pooling layer实现了网络对图像特征的多尺度特征识别。
2、上采样部分会融合特征提取部分的输出，这样做实际上是将多尺度特征融合在了一起，以最后一个上采样为例，它的特征既来自第一个卷积block的输出(同尺度特征)，也来自上采样的输出(大尺度特征)，这样的连接是贯穿整个网络的，你可以看到上图的网络中有四次融合过程，相对应的FCN网络只在最后一层进行融合。

Unet训练自己的数据

整个模型训练过程如下图所示：

一、Labelme制作标签

Labelme软件安装借鉴https://www.cnblogs.com/wangxiaocvpr/p/9997690.html，安装完成，打开图片后，界面如下：

点Edit->creat Polygons，然后绘制mask。Ctrl+S保存为JOSN文件，然后在目录下使用Prompt

输入labelme_json_to_dataset <文件名>.json，然后会在目录下产生一个文件夹，内容如下：

完成所有的标签工作后，将img和label文件各放一个文件夹。

二、Unet网络调参

网络调参涉及以下几个方面：
（1）加入BN层
（2）将最后一层激活函数替换成ReLU
（3）损失函数替换成mse
多分类一般最后一层原本是softmax，使用了这个激活函数跑完后，没有达到分割效果，所以替换成了之前做过的图对图项目激活函数，效果就出来了，纯属经验之谈，理论还没有进行验证。多分类的损失函数多是交叉熵，经过验证也是不能达到效果，替换成均方根误差。

三、训练与测试

此处可以参考开源代码https://github.com/haixiansheng/unet-keras-for-Multi-classification
项目中，训练测试单独分开，训练和测试数据以及标签都单独设立文件夹。以便于进行数据处理。

四、测试结果展示

项目要求分割出底座和胶水

分割电路板的胶水和底座

分割结果：绿色为胶水，红色为底座

五、总结

训练效果以及测试效果由于项目还没有完成，不便于展示，总体测试效果误差估计不到30像素。有意向的可以尝试用我们的网络进行分割实验。有所帮助的话记得关注公众号和star github项目。

往期文章

YOLO-Keras 跑通记录
使用Keras重现论文深度学习框架