Unet（pytorch）——制作自己的数据集（二分类）

 下方是pytorch Unet的源码 可以直接点进去下载

github源码

一、准备自己的数据集

在Pytorch-UNet-master\data下创建 imgs和masks两个文件夹

其中imgs存放rgb（.jpg格式)图像， masks存放标签（.png）图像

此时我手中已经有如下的jpg图像和png图像

jpg格式图像
​​​​​ png格式图像
 

但直接使用这样的png图像训练会报错

因为目标（白色）像素值是255  

所以我写了一个脚本可以批量把图像转化为像素值为0，1的图像

import os
import cv2
from PIL import Image

def binarization():
    # 获取目录下所有图片名
    filename = os.listdir(r"E:\..")
    print(filename)
    # os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表。
    base_dir = r"E:\.." # input
    new_dir  = r"E:\.."  # output
    for img in filename:
        name = img
        path1 = os.path.join(base_dir ,img)
        img = cv2.imread(path1)
        # print(img)
        Grayimg = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, thresh = cv2.threshold(Grayimg, 170, 1 ,cv2.THRESH_BINARY)
        cv2.imwrite('img.png', thresh)
        image = Image.open('img.png')
        # 有需要可对图像进行大小调整
        # image = image.resize((350, 350),Image.ANTIALIAS)
        path = os.path.join(new_dir ,name)
        image.save(path)

binarization()

之后生成的标签图为

调整后png图像

此时背景的像素值为0，目标像素为1

最后把调整过后的png放入masks文件夹，jpg放入imgs文件夹 

python train.py

即可开始训练