单纯为成功实现U-Net做个记录,并非保姆级教程。
本人代码:https://github.com/struggler176393/Unet_Pytorch
参考代码:https://github.com/qiaofengsheng/pytorch-UNet
环境
- pytorch 1.10.1
- python 3.7
- cuda11.3
- cudnn8.2.0
其实pytorch环境好像也就1.7开始有些函数名称变化有点大,要改改,所以这个代码估计1.7及更高版本的都能适用,1.7以前应该也可以。
还有一些常用的包比如numpy那些自行配置,我也不大记得了。
制作数据集
U-Net的数据集只需要JPG图片还有mask掩码就可以了,labelme太麻烦了,建议用EIseg,可以智能分割。
参考这篇blog:https://blog.csdn.net/qq_37541097/article/details/120154543
有个坑就是权重可能得另外找地方(官网)下载,blog里给的好像不行。
标注好后将JPG文件放入JPEGImages文件夹,mask掩码文件放入SegmentationClass文件夹。
运行
接下来挺简单了,想训练直接运行一下train.py,想测试直接运行一下test.py,然后输入预测图像的绝对路径。数据集小的话一两个小时训练的结果就挺可以了。 训练十分钟的效果,也还行了。