windows11下将Labelme标注数据转为YOLOV5训练数据集

完整代码：

import shutil
import os
import numpy as np
import json
from glob import glob
import cv2
from sklearn.model_selection import train_test_split
from utils.data_dir import root_dir


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def generate_mark_data(labelme_path, isUseTest = True):
    """

    Args:
        labelme_path: 标签路径
        isUseTest: 是否创建test集

    Returns:

    """
    # 获取待处理文件
    files = glob(labelme_path + "/*.json")
    files = [i.replace("\\", "/").split("/")[-1].split(".json")[0] for i in files]
    if isUseTest:
        trainval_files, test_files = train_test_split(files, test_size=0.1, random_state=55)
    else:
        trainval_files = files
        test_files = []  # 没有测试集
    train_files, val_files = train_test_split(trainval_files, test_size=0.1, random_state=55)
    return train_files, val_files, test_files


def ChangeToYolo5(yolov5_path, files, classes, txt_Name, img_suffix=".png"):
    tmp_path = f"{labelme_path}/tmp/"
    if not os.path.exists(tmp_path):  # 如果没有tmp文件夹则创建之
        os.makedirs(tmp_path)
    list_file = open(f'{tmp_path}%s.txt' % (txt_Name), 'w')
    for json_file_ in files:
        json_filename = labelme_path + "/" + json_file_ + ".json"
        imagePath = labelme_path + "/" + json_file_ + img_suffix
        list_file.write('%s/%s\n' % (yolov5_path, imagePath))
        out_file = open('%s/%s.txt' % (labelme_path, json_file_), 'w')
        json_file = json.load(open(json_filename, "r", encoding="utf-8"))
        height, width, channels = cv2.imread(labelme_path + "/" + json_file_ + img_suffix).shape
        for multi in json_file["shapes"]:
            points = np.array(multi["points"])
            xmin = min(points[:, 0]) if min(points[:, 0]) > 0 else 0
            xmax = max(points[:, 0]) if max(points[:, 0]) > 0 else 0
            ymin = min(points[:, 1]) if min(points[:, 1]) > 0 else 0
            ymax = max(points[:, 1]) if max(points[:, 1]) > 0 else 0
            label = multi["label"]
            if xmax <= xmin:
                pass
            elif ymax <= ymin:
                pass
            else:
                cls_id = classes.index(label)
                b = (float(xmin), float(xmax), float(ymin), float(ymax))
                bb = convert((width, height), b)
                out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
                print(json_filename, xmin, ymin, xmax, ymax, cls_id)


def make_voc_data(data_path, file_List, pathdepth=9):
    for file in file_List:
        if not os.path.exists(f'{data_path}/VOC/images/%s' % file):
            os.makedirs(f'{data_path}/VOC/images/%s' % file)
        if not os.path.exists(f'{data_path}/VOC/labels/%s' % file):
            os.makedirs(f'{data_path}/VOC/labels/%s' % file)
        f = open(f'{data_path}/LabelmeData/img/8P203/tmp/%s.txt' % file, 'r')
        lines = f.readlines()
        for line in lines:
            line = "/".join(line.split('/')[-pathdepth:]).strip()
            shutil.copy(line, f"{data_path}/VOC/images/%s" % file)
            line = line.replace('jpg', 'txt')
            shutil.copy(line, f"{data_path}/VOC/labels/%s/" % file)


if __name__ == '__main__':
    labelme_path = f"{root_dir}/data/LabelmeData/img/8P203"
    classes = ["1", "2", "3", "4", "5", "6"]
    train_files, val_files, test_files = generate_mark_data(labelme_path)
    yolov5_path = labelme_path
    ChangeToYolo5(yolov5_path, train_files, classes, "train")
    ChangeToYolo5(yolov5_path, val_files, classes, "val")
    ChangeToYolo5(yolov5_path, test_files, classes, "test")
    data_path = f"{root_dir}/data"
    file_List = ["train", "val", "test"]
    make_voc_data(data_path, file_List)
    '''
    file1 = open("tmp/train.txt", "r")
    file2 = open("tmp/val.txt", "r")
    file_list1 = file1.readlines()  # 将所有变量读入列表file_list1
    file_list2 = file2.readlines()  # 将所有变量读入列表file_list2
    file3 = open("tmp/trainval.txt", "w")
    for line in file_list1:
        print(line)
        file3.write(line)
    for line in file_list2:
        print(line)
        file3.write(line)
    '''

其中的路径root_dir来源如下：