# -*- coding: utf-8 -*- import os import numpy as np import json from glob import glob import cv2 import shutil import yaml from sklearn.model_selection import train_test_split from tqdm import tqdm # 获取当前路径 ROOT_DIR = os.getcwd() ''' 统一图像格式 ''' def change_image_format(label_path=ROOT_DIR, suffix='.png'): """ 统一当前文件夹下所有图像的格式，如'.jpg' :param suffix: 图像文件后缀 :param label_path:当前文件路径 :return: """ externs = ['png', 'jpg', 'JPEG', 'BMP', 'bmp'] files = list() # 获取尾缀在ecterns中的所有图像 for extern in externs: files.extend(glob(label_path + "\\*." + extern)) # 遍历所有图像，转换图像格式 for file in files: name = ''.join(file.split('.')[:-1]) file_suffix = file.split('.')[-1] if file_suffix != suffix.split('.')[-1]: # 重命名为jpg new_name = name + suffix # 读取图像 image = cv2.imread(file) # 重新存图为jpg格式 cv2.imwrite(new_name, image) # 删除旧图像 os.remove(file) ''' 读取所有json文件，获取所有的类别 ''' def get_all_class(file_list, label_path=ROOT_DIR): """ 从json文件中获取当前数据的所有类别 :param file_list:当前路径下的所有文件名 :param label_path:当前文件路径 :return: """ # 初始化类别列表 classes = list() # 遍历所有json,读取shape中的label值内容，添加到classes for filename in tqdm(file_list): json_path = os.path.join(label_path, filename + '.json') json_file = json.load(open(json_path, "r", encoding="utf-8")) for item in json_file["shapes"]: label_class = item['label'] if label_class not in classes: classes.append(label_class) print('read file done') return classes ''' 划分训练集、验证机、测试集 ''' def split_dataset(label_path, test_size=0.3, isUseTest=False, useNumpyShuffle=False): """ 将文件分为训练集，测试集和验证集 :param useNumpyShuffle: 使用numpy方法分割数据集 :param test_size: 分割测试集或验证集的比例 :param isUseTest: 是否使用测试集，默认为False :param label_path:当前文件路径 :return: """ # 获取所有json files = glob(label_path + "\\*.json") files = [i.replace("\\", "/").split("/")[-1].split(".json")[0] for i in files] if useNumpyShuffle: file_length = len(files) index = np.arange(file_length) np.random.seed(32) np.random.shuffle(index) # 随机划分 test_files = None # 是否有测试集 if isUseTest: trainval_files, test_files = np.array(files)[index[:int(file_length * (1 - test_size))]], np.array(files)[ index[int(file_length * (1 - test_size)):]] else: trainval_files = files # 划分训练集和测试集 train_files, val_files = np.array(trainval_files)[index[:int(len(trainval_files) * (1 - test_size))]], \ np.array(trainval_files)[index[int(len(trainval_files) * (1 - test_size)):]] else: test_files = None if isUseTest: trainval_files, test_files = train_test_split(files, test_size=test_size, random_state=55) else: trainval_files = files train_files, val_files = train_test_split(trainval_files, test_size=test_size, random_state=55) return train_files, val_files, test_files, files ''' 生成yolov5的训练、验证、测试集的文件夹 ''' def create_save_file(label_path=ROOT_DIR): """ 按照训练时的图像和标注路径创建文件夹 :param label_path:当前文件路径 :return: """ # 生成训练集 train_image = os.path.join(label_path, 'train', 'images') if not os.path.exists(train_image): os.makedirs(train_image) train_label = os.path.join(label_path, 'train', 'labels') if not os.path.exists(train_label): os.makedirs(train_label) # 生成验证集 val_image = os.path.join(label_path, 'valid', 'images') if not os.path.exists(val_image): os.makedirs(val_image) val_label = os.path.join(label_path, 'valid', 'labels') if not os.path.exists(val_label): os.makedirs(val_label) # 生成测试集 test_image = os.path.join(label_path, 'test', 'images') if not os.path.exists(test_image): os.makedirs(test_image) test_label = os.path.join(label_path, 'test', 'labels') if not os.path.exists(test_label): os.makedirs(test_label) return train_image, train_label, val_image, val_label, test_image, test_label ''' 转换，根据图像大小，返回box框的中点和高宽信息 ''' 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 push_into_file(file, images, labels, label_path=ROOT_DIR, suffix='.jpg'): """ 最终生成在当前文件夹下的所有文件按image和label分别存在到训练集/验证集/测试集路径的文件夹下 :param file: 文件名列表 :param images: 存放images的路径 :param labels: 存放labels的路径 :param label_path: 当前文件路径 :param suffix: 图像文件后缀 :return: """ # 遍历所有文件 for filename in file: # 图像文件 image_file = os.path.join(label_path, filename + suffix) # 标注文件 label_file = os.path.join(label_path, filename + '.txt') # yolov5存放图像文件夹 if not os.path.exists(os.path.join(images, filename + suffix)): try: shutil.move(image_file, images) except OSError: pass # yolov5存放标注文件夹 if not os.path.exists(os.path.join(labels, filename + suffix)): try: shutil.move(label_file, labels) except OSError: pass ''' ''' def json2txt(classes, txt_Name='allfiles', label_path=ROOT_DIR, suffix='.png'): """ 将json文件转化为txt文件，并将json文件存放到指定文件夹 :param classes: 类别名 :param txt_Name:txt文件，用来存放所有文件的路径 :param label_path:当前文件路径 :param suffix:图像文件后缀 :return: """ store_json = os.path.join(label_path, 'json') if not os.path.exists(store_json): os.makedirs(store_json) _, _, _, files = split_dataset(label_path) if not os.path.exists(os.path.join(label_path, 'tmp')): os.makedirs(os.path.join(label_path, 'tmp')) list_file = open('tmp/%s.txt' % txt_Name, 'w') for json_file_ in tqdm(files): json_filename = os.path.join(label_path, json_file_ + ".json") imagePath = os.path.join(label_path, json_file_ + suffix) list_file.write('%s\n' % imagePath) out_file = open('%s/%s.txt' % (label_path, json_file_), 'w') json_file = json.load(open(json_filename, "r", encoding="utf-8")) if os.path.exists(imagePath): height, width, channels = cv2.imread(imagePath).shape for multi in json_file["shapes"]: if len(multi["points"][0]) == 0: out_file.write('') continue 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) if not os.path.exists(os.path.join(store_json, json_file_ + '.json')): try: shutil.move(json_filename, store_json) except OSError: pass ''' 创建yaml文件 ''' def create_yaml(classes, label_path, isUseTest=False): nc = len(classes) if not isUseTest: desired_caps = { 'path': label_path, 'train': 'train/images', 'val': 'valid/images', 'nc': nc, 'names': classes } else: desired_caps = { 'path': label_path, 'train': 'train/images', 'val': 'valid/images', 'test': 'test/images', 'nc': nc, 'names': classes } yamlpath = os.path.join(label_path, "data" + ".yaml") # 写入到yaml文件 with open(yamlpath, "w+", encoding="utf-8") as f: for key, val in desired_caps.items(): yaml.dump({key: val}, f, default_flow_style=False) # 首先确保当前文件夹下的所有图片统一后缀，如.jpg，如果为其他后缀，将suffix改为对应的后缀，如.png def ChangeToYolo5(label_path=r"D:\storydata", suffix='.png', test_size=0.1, isUseTest=False): """ 生成最终标准格式的文件 :param test_size: 分割测试集或验证集的比例 :param label_path:当前文件路径 :param suffix: 文件后缀名 :param isUseTest: 是否使用测试集 :return: """ # step1:统一图像格式 change_image_format(label_path) # step2:根据json文件划分训练集、验证集、测试集 train_files, val_files, test_file, files = split_dataset(label_path, test_size=test_size, isUseTest=isUseTest) # step3:根据json文件，获取所有类别 classes = get_all_class(files) # step4:将json文件转化为txt文件，并将json文件存放到指定文件夹 json2txt(classes) # step5:创建yolov5训练所需的yaml文件 create_yaml(classes, label_path, isUseTest=isUseTest) # step6:生成yolov5的训练、验证、测试集的文件夹 train_image, train_label, val_image, val_label, test_image, test_label = create_save_file(label_path) # step7:将所有图像和标注文件，移动到对应的训练集、验证集、测试集 push_into_file(train_files, train_image, train_label, suffix=suffix) # 将文件移动到训练集文件中 push_into_file(val_files, val_image, val_label, suffix=suffix) # 将文件移动到验证集文件夹中 if test_file is not None: # 如果测试集存在，则将文件移动到测试集文件中 push_into_file(test_file, test_image, test_label, suffix=suffix) print('create dataset done') if __name__ == "__main__": ChangeToYolo5()