yolov5训练自已的数据集该如何处理

 1.划分数据集

def spilt_data():
    '''
    拆分数据集: 训练集:验证集:测试集 = 9:1
    '''
    import os
    import random
    import glob
    train_val_percent = 1
    train_percent = 0.9  #这里的 train_percent 是指占 train_val_percent 中的
    # xml 标签文件
    base_path = './data/datasets'
    xml_file_path = base_path + '/*.json'
    txt_save_path = base_path + '/txts'
    # 判断txt_save_path文件夹路径如果不存在,就创建出来
    if not os.path.exists(txt_save_path):
        os.makedirs(txt_save_path)
    # 统计标签文件json的总数
    total_json = glob.glob(xml_file_path)
    num = len(total_json)
    list_index = range(num)  # 获取随机0 到 num-1的数字,作为索引值
    # 计算训练集和验证集的数据量
    train_val_num = int(num * train_val_percent)
    # 计算训练集的数据量
    train_num = int(train_val_num * train_percent)
    # 获取列表 list_index 中 指定长度的随机数
    train_val_indexs = random.sample(list_index, train_val_num)
    train_indexs = random.sample(train_val_indexs, train_num)
    # 打开四个文件 trainval.txt、test.txt、train.txt、val.txt 用来保存分出来的数据名称
    file_train_val = open(os.path.join(txt_save_path, 'trainval.txt'), 'w')
    file_train = open(os.path.join(txt_save_path, 'train.txt'), 'w')
    file_test = open(os.path.join(txt_save_path, 'test.txt'), 'w')
    file_val = open(os.path.join(txt_save_path, 'val.txt'), 'w')
    for i in list_index:
        # 获取到xml的文件名称,不要后缀
        name = total_json[i].split('\\')[-1]
        name = os.path.splitext(name)[0] + '\n'
        # 判断当前的索引值i是在哪部分
        if i in train_val_indexs:
            # 在训练集或者验证集中
            file_train_val.write(name)
            if i in train_indexs:
                # 在训练集中
                file_train.write(name)
            else:
                # 在验证集中
                file_val.write(name)
        else:
            # 在测试集中
            file_test.write(name)
 
    # 关闭资源
    file_train_val.close()
    file_train.close()
    file_test.close()
    file_val.close()
    
if __name__== "__main__":
    spilt_data()

2.转为yolo格式

import json
import os.path
import shutil
from tqdm import tqdm
import cv2
 
sets = ['train', 'test','val']
#这里使用要改成自己的类别
classes = ['','']
# 解析单个json文件
def json2yolo(image_id):
    pass
    in_file_path = json_path % image_id
    img_file_path = images_jpg_path % image_id
    out_file = open(labels_path + '/%s.txt' % image_id, 'w', encoding='utf-8')
    with open(in_file_path, 'r') as file:
        data = json.load(file)
 
    # 获取图片的尺寸
    img = cv2.imread(img_file_path)
    img_h, img_w, c = img.shape
    
    # 遍历所有标注的目标
    list_data = data['shapes']
    for obj in list_data:
        # 获取类别
        cls = obj['label']
        # 判断找到的类别是否是我们需要的
        if cls not in classes:
            continue
        # 获取类别id
        cls_id = classes.index(cls)
        # print(cls_id)
        points = obj['points']
        xmin = int(points[0][0])
        ymin = int(points[0][1])
        xmax = int(points[1][0])
        ymax = int(points[1][1])
        # print(xmin, ymin, xmax, ymax)
        # exit()
        # 标注越界修正
        if xmax > img_w:
            xmax = img_w
        if ymax > img_h:
            ymax = img_h
        # print(cls_id, xmin, ymin, xmax, ymax)
 
        # 转换成 中心点 和 宽高 的形式
        w = xmax - xmin
        h = ymax - ymin
        cx = xmin + w / 2
        cy = ymin + h / 2
 
        # 归一化 , 保留6位小数
        w = round(w / img_w, 6)
        h = round(h / img_h, 6)
        cx = round(cx / img_w, 6)
        cy = round(cy / img_h, 6)
        # print(cls_id, cx, cy, w, h)
        # 结果保存到数据labels文件夹中的txt文件
        out_file.write(str(cls_id) + " " + str(cx) + " " + str(cy) +
                       " " + str(w) + " " + str(h) + '\n')
 
    # 释放资源
    out_file.close()
 
 
def read_files():
    # 获取 train.txt, test.txt ,val.txt 文件中的内容
    for img_set in sets:
        # print(img_set)  # train \ test \ val
        img_path = imageSets_path % (img_set)  # 找到 ImageSets文件夹下面的txt文件
        # 打开train.txt文件
        with open(img_path, 'r', encoding='utf-8') as file:
            # 获取里面所有的图片名称
            image_ids = [item.strip() for item in file.readlines()]
            # print(image_ids)
 
        # 将获取的 image_id 对应的图片路径保存到txt中
        list_file = open(base_path + '/%s.txt' % (img_set), 'w')
        # 遍历图片名称
        for image_id in tqdm(image_ids):  # 可以看见执行的进度
            # print(image_id)
            json2yolo(image_id)
            list_file.write(images_path + '/%s.jpg\n' % (image_id))
        list_file.close()
 
 
base_path = '/data/coding/yolov5-master/data/datasets'
imageSets_path = base_path + '/txts/%s.txt'
json_path = base_path + '/%s.json'
images_jpg_path = base_path + '/%s.jpg'
labels_path = base_path + '/create_datas/labels'  # 存放转换后所有的txt文件
images_path = base_path + '/create_datas/images'
 
if __name__ == '__main__':
    print("数据处理 开始")
    # 判断保存的文件夹不存在就创建
    if not os.path.exists(labels_path):
        os.makedirs(labels_path)
    if not os.path.exists(images_path):
        os.makedirs(images_path)
    # 将 JPEGImages 文件夹下面所有的图片保存到 images文件夹下面
    for filename in os.listdir(base_path):
        if filename.endswith('.jpg'):
            # 构建源文件和目标文件的完整路径
            source_file = os.path.join(base_path, filename)
            target_file = os.path.join(images_path, filename)
            # 复制文件
            shutil.copy2(source_file, target_file)  
 
    read_files()
    print("数据处理 完成")

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值