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

版权声明:本文为CSDN博主「无穷升高的卡农」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/weixin_40802676/article/details/115479672

我还没有学会写个人说明!

暂无评论

发表评论

相关推荐

使用PyTorch实现目标检测与跟踪

欢迎关注 “小白玩转Python”,发现更多 “有趣”引言在昨天的文章中,我们介绍了如何在PyTorch中使用您自己的图像来训练图像分类器,然后使用它来进行图像识别。本文将展示如何使用预训练的分类器检测

实战 | 实时的目标检测与识别简单应用

吃粽子迎端午计算机视觉研究院专栏作者:Edison_G最近总是有很多入门的朋友问我,我进入计算机视觉这个领域难不难?是不是要学习很多知识?到底哪个方向比较好?长按扫描二维码关