利用U-Net语义分割进行自己训练集的训练和预测

利用U-Net语义分割进行自己训练集的训练和预测

引言

利用U-Net语义分割网络进行三类图像分割（含背景），总体步骤分为以下三步：
（1）制作自己的数据集（杯子数据集）
通过labelme进行图像打标（注意：最好利用低版本，高版本有Bug）。样图如下：

我是用来分割出保温杯和茶杯，值得一提的是背景在进行训练时也算作一类。
（2）转换json文件
打标结束后，会出现json文件，通过json_to_dataset.py将json文件转换为原图和标签图，标签图如下：

当然，打标过程是手动的！
（3）训练数据集
我采用的是vgg16的预训练权重，评价指标采用交叉熵验证和dice_loss，根据最终的Loss判断训练效果的好坏。

（4）最终也是训练好了网络（.pth），通过替换这个模型，利用predict.py来进行预测，预测效果如下，效果还是非常好的！