YOLOV5模型训练之JSON格式的数据集制作，一个py程序搞，搞不定邮件联系，我帮你！！！！

# -*- coding: utf-8 -*-
"""
数据集转换脚本，包含以下内容
1、json文件转txt文件
2、对数据集进行划分，并存放至相应目录
注意：
1、将此程序放置yolov5工程目录下的data文件夹下
2、图像和标签文件夹的名字和我的文件名称改为一致
"""
import json
import os
import cv2
import shutil
import random


# ！！！！这里是重点，需要查看自己的数据集名称对不对，不对的话，把自己的文件夹重命名一下
json_file_path = './Annotations/'                  # 检查和自己的json文件夹名称是否一致
img_folder_path = './SourceImages/'        # 检查和自己的图像文件夹名称是否一致

# 将自己数据集的类别对应写一下，几类就写几个,一定要从0开始
label_mapping = {
    "nicaishenmeyisi":0,
    "jiushiwaner": 1,
}


# 数据集划分比例，训练集70%，验证集20%， 测试集10%
train_percent = 0.7
val_percent = 0.2
test_percent = 0.1

if not os.path.exists('./temp_labels/'):      # 临时标签文件夹位置，这不要改动，之后会删掉的
    os.makedirs('./temp_labels/')
txt_folder_path = './temp_labels/'


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

          
#保存为绝对坐标形式 :label x1 y1 x2 y2
def absolute_coordinate_txt(img_name,json_d,img_path):
    src_img=cv2.imread(img_path)
    h,w = src_img.shape[:2]
    txt_name = img_name.split(".")[0]+".txt"
    txt_path = os.path.join(txt_folder_path,txt_name)
    # print("txt_path:\t",txt_path)
    with open(txt_path,'w') as f:
        for item in json_d["shapes"]:
            point=item['points']
            x1 = point[0][0]
            y1 = point[0][1]
            x2 = point[1][0]
            y2 = point[1][1]
            b = (float(x1), float(x2), float(y1), float(y2))
            bb = convert((w, h), b)
            f.write(str(label_mapping[item['label']]) + " " + " ".join([str(a) for a in bb]) + '\n')


# 如果json标签有问题，print输出一下，可以查看哪个json文件有问题
for jsonfile in os.listdir(json_file_path):
    temp_path=os.path.join(json_file_path,jsonfile)
    # print("json_path:\t",temp_path) 
    jsonfile_path=temp_path
    with open(jsonfile_path, "r", encoding='utf-8') as f:
        json_d = json.load(f)    
        img_name =  jsonfile[:-4]+"jpg"
        img_path=os.path.join(img_folder_path,img_name)
        # print("img_path:\t",img_path)
        absolute_coordinate_txt(img_name,json_d,img_path)
# *********************************************** #
#  parent folder
#  ----images
#       ----train
#       ----val
#       ----test
#  ----labels
#       ----train
#       ----val
#       ----test
def create_dir():
    if not os.path.exists('./images/'):
        os.makedirs('./images/')
    if not os.path.exists('./labels/'):
        os.makedirs('./labels/')

    if not os.path.exists('./images/train/'):
        os.makedirs('./images/train')
    if not os.path.exists('./images/val/'):
        os.makedirs('./images/val/')
    if not os.path.exists('./images/test/'):
        os.makedirs('./images/test/')

    if not os.path.exists('./labels/train/'):
        os.makedirs('./labels/train/')
    if not os.path.exists('./labels/val/'):
        os.makedirs('./labels/val/')
    if not os.path.exists('./labels/test/'):
        os.makedirs('./labels/test/')
    return


# *********************************************** #
print("************************************************************************")
# 读取所有的txt文件
create_dir()
total_txt = os.listdir(txt_folder_path)
num_txt = len(total_txt)
list_all_txt = range(num_txt)  # 范围 range(0, num)

num_train = int(num_txt * train_percent)
num_val = int(num_txt * val_percent)
num_test = num_txt - num_train - num_val

train = random.sample(list_all_txt, num_train)
# train从list_all_txt取出num_train个元素
# 所以list_all_txt列表只剩下了这些元素：val_test
val_test = [i for i in list_all_txt if not i in train]
# 再从val_test取出num_val个元素，val_test剩下的元素就是test
val = random.sample(val_test, num_val)
# 检查两个列表元素是否有重合的元素
# set_c = set(val_test) & set(val)
# list_c = list(set_c)
# print(list_c)
# print(len(list_c))

print("训练集数目：{}, 验证集数目：{},测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
for i in list_all_txt:
    name = total_txt[i][:-4]

    srcImage = img_folder_path + name + '.jpg'
    srcLabel = txt_folder_path + name + '.txt'

    if i in train:
        dst_train_Image = './images/train/' + name + '.jpg'
        dst_train_Label = './labels/train/' + name + '.txt'
        shutil.copyfile(srcImage, dst_train_Image)
        shutil.copyfile(srcLabel, dst_train_Label)
    elif i in val:
        dst_val_Image = './images/val/' + name + '.jpg'
        dst_val_Label = './labels/val/' + name + '.txt'
        shutil.copyfile(srcImage, dst_val_Image)
        shutil.copyfile(srcLabel, dst_val_Label)
    else:
        dst_test_Image = './images/test/' + name + '.jpg'
        dst_test_Label = './labels/test/' + name + '.txt'
        shutil.copyfile(srcImage, dst_test_Image)
        shutil.copyfile(srcLabel, dst_test_Label)
shutil.rmtree(txt_folder_path)