U-Net训练自己的数据集

本文详细介绍了如何使用U-Net模型训练自己的数据集,包括数据准备、下载Pytorch-UNet存储库、设置训练参数并进行训练,以及最后的模型测试。数据准备阶段涉及将labelme的json文件转换为二值化掩膜。训练过程简单,只需将图片和掩膜放入相应文件夹,并通过pythontrain.py脚本启动。测试阶段使用predict.py脚本对模型进行预测。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

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
评论 14
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值