U-Net训练自己的数据集
U-Net训练自己的数据集1:数据准备1:数据准备标注部分不再介绍。
·
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
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐



所有评论(0)