U-Net训练自己的数据集

1:数据准备

我说使用的数据是labelme制作的。json文件保存的是对应图片中所有目标的边界点坐标。

但是UNet训练却使用的是原始图像及其对应的二值化掩膜。就像下面这样:
在这里插入图片描述
在这里插入图片描述
所以需要把labelme输出的json文件绘制出二值化掩膜:方法可参考:已知分割数据集图像的多边形边界,绘制其二值化掩膜
注意要是单通道的二值化掩膜。

2:下载存储库

Pytorch-UNet
代码并不需要修改什么,直接就可以运行了。只要把数据集放到正确的位置:
在这里插入图片描述
imgs里存放原始图片
masks里存放二值化掩膜图片

3:训练

可以自己设置一些参数,比如epochs。

python train.py --epochs=100

在这里插入图片描述
每个epoch结束,会保存这次训练之后的模型。存在一个checkpoints文件夹里

4:测试

python predict.py -i test/image.jpg -o output.jpg --model=checkpoints/CP_epoch1.pth

版权声明:本文为weixin_40802676原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接和本声明。