Unet训练自己的数据集保姆级教程

1.利用labelme进行数据标注

1.1Labelme 安装方法

首先安装 Anaconda，然后运行下列命令：

##################
## for Python 2 ##
##################
conda create --name=labelme python=2.7
source activate labelme
# conda install -c conda-forge pyside2
conda install pyqt
pip install labelme
# 如果想安装最新版本，请使用下列命令安装：
# pip install git+https://github.com/wkentaro/labelme.git
 
##################
## for Python 3 ##
##################
conda create --name=labelme python=3.6
source activate labelme
# conda install -c conda-forge pyside2
# conda install pyqt
pip install pyqt5  # pyqt5 can be installed via pip on python3
pip install labelme
输入以下指令打开
labelme

1.2Labelme 使用教程

使用 labelme 进行场景分割标注的教程详见：labelme

2.将JSON格式标注文件转换为png格式

新建文件夹如图所示：JPEGImages存放图片，SegmentationClass_Origin存放json标签

运行如下代码进行格式转换，注意填对文件地址和种类数量

import base64
import json
import os
import os.path as osp

import numpy as np
import PIL.Image
from labelme import utils

'''
制作自己的语义分割数据集需要注意以下几点：
1、我使用的labelme版本是3.16.7，建议使用该版本的labelme，有些版本的labelme会发生错误，
   具体错误为：Too many dimensions: 3 > 2
   安装方式为命令行pip install labelme==3.16.7
2、此处生成的标签图是8位图，此时每个像素点的值就是这个像素点所属的种类。
'''
if __name__ == '__main__':
    jpgs_path   = "datasets/JPEGImages"
    pngs_path   = "datasets/SegmentationClass"
    classes     = ["_background_","1"]  ##种类填什么名称都可以
    # classes     = ["_background_","cat","dog"]
    
    count = os.listdir("./datasets/SegmentationClass_Origin/")
    for i in range(0, len(count)):
        path = os.path.join("./datasets/SegmentationClass_Origin", count[i])

        if os.path.isfile(path) and path.endswith('json'):
            data = json.load(open(path))
            
            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')

            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {'_background_': 0}
            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value
            
            # label_values must be dense
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))
            
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
            

           # 保存处理后的原始图像为 .jpg 文件
            PIL.Image.fromarray(img).save(osp.join(jpgs_path, count[i].split(".")[0]+'.jpg'))

             # 创建标签图像，并将每个类别映射为对应的整数值
            new = np.zeros([np.shape(img)[0], np.shape(img)[1]], dtype=np.uint8)
            for name in label_names:
                index_json = label_names.index(name)
                index_all = classes.index(name)
                new = new + (index_all * (lbl == index_json)).astype(np.uint8)

           # 保存标签图像为灰度图的 .png 文件
            #utils.lblsave(osp.join(pngs_path, count[i].split(".")[0]+'.png'), new)
            PIL.Image.fromarray(new, mode='L').save(osp.join(pngs_path, count[i].split(".")[0] + '.png'))
            print('Saved ' + count[i].split(".")[0] + '.jpg and ' + count[i].split(".")[0] + '.png')

转换完成后可以通过data_process/check_label.py文件查看标签的值，比如有两个种类，那么png标签的值就只能是0和1。

3.将标签和图片转化为统一大小

import os
import cv2

def resize_images(input_folder1, input_folder2, output_folder1, output_folder2, size=(640, 640)):
    # 确保输出文件夹存在
    os.makedirs(output_folder1, exist_ok=True)
    os.makedirs(output_folder2, exist_ok=True)

    # 获取文件夹1中所有.jpg图片
    jpg_files = [f for f in os.listdir(input_folder1) if f.endswith('.jpg')]

    # 获取文件夹2中所有.png图片
    png_files = [f for f in os.listdir(input_folder2) if f.endswith('.png')]

    # 遍历.jpg文件并处理
    for jpg_file in jpg_files:
        # 构造完整的文件路径
        jpg_path = os.path.join(input_folder1, jpg_file)
        # 读取.jpg图片为灰度图
        img_jpg = cv2.imread(jpg_path)
        # 调整大小
        img_jpg_resized = cv2.resize(img_jpg, size)
        # 保存到新的文件夹
        cv2.imwrite(os.path.join(output_folder1, jpg_file), img_jpg_resized)

        # 处理对应的.png图片
        png_file = jpg_file.replace('.jpg', '.png')
        if png_file in png_files:
            png_path = os.path.join(input_folder2, png_file)
            # 读取.png图片为灰度图
            img_png = cv2.imread(png_path, cv2.IMREAD_GRAYSCALE)
            # 调整大小
            img_png_resized = cv2.resize(img_png, size)
            # 保存到新的文件夹
            cv2.imwrite(os.path.join(output_folder2, png_file), img_png_resized)

# 定义输入输出文件夹路径
input_folder1 = 'F:/Improve/unet_improve/unet/datasets/JPEGImages'
input_folder2 = 'F:/Improve/unet_improve/unet/datasets/SegmentationClass'
output_folder1 = 'F:/Improve/unet_improve/unet/datasets/JPEGImages2'
output_folder2 = 'F:/Improve/unet_improve/unet/datasets/SegmentationClass2'

# 调用函数进行图片处理
resize_images(input_folder1, input_folder2, output_folder1, output_folder2, size=(640, 640))

4.划分数据集

按照不同的比例将数据集划分为训练验证集和测试集

import os
import shutil
import random


def split_dataset(image_folder, label_folder, train_ratio, train_image_folder, train_label_folder, test_image_folder,
                  test_label_folder):
    # 获取所有图片文件名
    image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]

    # 将文件顺序打乱
    random.shuffle(image_files)

    # 计算训练集的数量
    train_size = int(len(image_files) * train_ratio)

    # 创建输出目录
    os.makedirs(train_image_folder, exist_ok=True)
    os.makedirs(train_label_folder, exist_ok=True)
    os.makedirs(test_image_folder, exist_ok=True)
    os.makedirs(test_label_folder, exist_ok=True)

    # 划分训练集和测试集
    train_files = image_files[:train_size]
    test_files = image_files[train_size:]

    # 复制文件到训练集文件夹
    for file in train_files:
        label_file = file.replace('.jpg', '.png')
        shutil.copy(os.path.join(image_folder, file), os.path.join(train_image_folder, file))
        shutil.copy(os.path.join(label_folder, label_file), os.path.join(train_label_folder, label_file))

    # 复制文件到测试集文件夹
    for file in test_files:
        label_file = file.replace('.jpg', '.png')
        shutil.copy(os.path.join(image_folder, file), os.path.join(test_image_folder, file))
        shutil.copy(os.path.join(label_folder, label_file), os.path.join(test_label_folder, label_file))


# 示例使用
image_folder = 'F:/Improve/unet_improve/unet/640voc/images'
label_folder = 'F:/Improve/unet_improve/unet/640voc/i-masks'
train_image_folder = 'F:/Improve/unet_improve/unet/640voc/train_images'
train_label_folder = 'F:/Improve/unet_improve/unet/640voc/train_masks'
test_image_folder = 'F:/Improve/unet_improve/unet/640voc/test_images'
test_label_folder = 'F:/Improve/unet_improve/unet/640voc/test_masks'

split_dataset(image_folder, label_folder, train_ratio=0.8,
              train_image_folder=train_image_folder,
              train_label_folder=train_label_folder,
              test_image_folder=test_image_folder,
              test_label_folder=test_label_folder)

5.训练

使用命令训练：
python train.py 
--model unet 
--imgs ../train_data_aug/imgs/ 
--masks ../train_data_aug/i-masks/ 
--save ../i-checkpoints/ 
--epochs 50 
--batch-size 64 
--scale 0.5 
-w 1e-4 
-epc 5 
-ls dice+ce 
-o adam 
--amp

或者在train.py中直接修改相应的参数，然后右键运行即可

参数说明：

--scale 0.5 表示将原始图片缩放到0.5倍大小进行训练，比如原始图片大小为640X640，这里填0.5，就表示训练图像大小为320X320

6.评估

在eval_test.py中添加pth权重路径并修改种类数classes,然后运行以下命令即可得到评估参数

其中input为图片地址 output为标签地址
python eval_test.py --input 640voc/images/ --output 640voc/i-masks

7.预测

同样修改权重路径和种类即可，然后运行下面的命令
python predict.py --input /home/cdg/test/unet-pytorch-main/VOCdevkit2/VOC2007/JPEGImages2/00000.jpg --viz --no-save

效果图：

整理不易，完整代码收取少许费用

完整代码地址：unet