U-net复现pytorch版本 以及制作自己的数据集并训练

GitHub链接：

GitHub – milesial/Pytorch-UNet: PyTorch implementation of the U-Net for image semantic segmentation with high quality images

论文链接：

https://arxiv.org/pdf/1505.04597v1.pdf

复现：

1.从GitHub下载源码以后，配置U-net的环境。

matplotlib
numpy
Pillow
torch
torchvision
tensorboard
future
tqdm

2.配置完成后进行训练，但是我们没有数据集呀，这里我找到了一个医学细胞壁的数据集，只有30张灰度照片，各位老铁可以用这个进行训练~

百度网盘链接

链接：链接：https://pan.baidu.com/s/1EOuJzgs8UyPSCfJsxBDFuA 
提取码：jt55 
 

然后将其放在U-net master /data 下面 ，修改train.py文件的路径

当然因为这个细胞是灰度图，所以要修改channel=1，

然后在部署好的环境里运行：

python train.py

然后就会看到如下：

这样的话说明正在训练咯，最终的dice 系数会接近1，那说明训练效果还不错。

3.利用训练好的模型进行训练，我们可以在U-net master /checkpoints路径下看到保存好的模型

所以我们修改U-net master /predict.py 文件中模型的路径：

然后我们进行预测：

python predict.py -i image.jpg -o output.jpg

iamge.jpg是输入 output.jpg是预测的结果，如图：

到此，我们复现完成。

制作自己的数据集：

工具：labelme + 自己获取的照片 

1.首先要将照片批量处理成512*512像素的照片哈

2.开始标注：

安装labelme:可以直接在anaconda下运行pip install labelme即安装完成

然后在命令行输入labelme即可运行，

labelme的使用教程网上一大把，选中你想识别的类标注即可。

3.标注完得到的json文件不能直接用来训练，要批量转化成png才能用作U-net的训练集，

至于怎么转化成png可以参考我的另一篇博客：

labelme批量实现json_to_dataset方法（2021）_candice5566的博客-CSDN博客

转化完成之后，就得到如下图：

mask-png 中就是标注好的照片了。

然后将image和mask改成相同的名字，至于怎么批量转化：

import os
import sys

# 修改生成的文件名字
def rename():
    path = input("请输入路径(例如D:\\\\picture)：")
    name = input("请输入开头名:")
    startNumber = input("请输入开始数:")
    fileType= input("请输入后缀名（如 .jpg、.txt等等）:")
    print("正在生成以"+name+startNumber+fileType+"迭代的文件名")
    count = 0
    filelist = os.listdir(path)
    for files in filelist:
        Olddir = os.path.join(path,files)
        if os.path.isdir(Olddir):
            continue
        Newdir = os.path.join(path, name+str(count+int(startNumber))+fileType)
        os.rename(Olddir,Newdir)
        count += 1
    print("一共修改了"+str(count)+"个文件")

rename()

所以我们的数据集就有了。把这个拷贝到我们项目的data下面，然后和复现一样，进行训练就可以了哈。

记得修改channel，类别，之类的哈。

over，完结撒花~

————————-2021.11.14———-

重要！！！！重新复现了一下U-Net,评论区常见的问题我整理到一篇新的文章中，大家可以去看看解决方法：

https://blog.csdn.net/candice5566/article/details/121318522

来源：奶盖芒果