赞
踩
path: D:\BaiduPan\labelmeCatAndDog
train: images/train
val: images/val
names:
0: dog
1: cat
# -*- 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 from PIL import Image ''' 统一图像格式 ''' def change_image_format(label_path, suffix='.jpg'): """ 统一当前文件夹下所有图像的格式,如'.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 index,file in enumerate(tqdm(files)): name = ''.join(file.split('.')[:-1]) file_suffix = file.split('.')[-1] if file_suffix != suffix.split('.')[-1]: # 重命名为jpg new_name = name + suffix # 读取图像 image = Image.open(file) image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR) # 重新存图为jpg格式 cv2.imwrite(new_name, image) # 删除旧图像 os.remove(file) ''' 读取所有json文件,获取所有的类别 ''' def get_all_class(file_list, label_path): """ 从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(ROOT_DIR): print('step6:生成yolov5的训练、验证、测试集的文件夹') # 生成训练集 train_image = os.path.join(ROOT_DIR, 'images','train') if not os.path.exists(train_image): os.makedirs(train_image) train_label = os.path.join(ROOT_DIR, 'labels','train') if not os.path.exists(train_label): os.makedirs(train_label) # 生成验证集 val_image = os.path.join(ROOT_DIR, 'images', 'val') if not os.path.exists(val_image): os.makedirs(val_image) val_label = os.path.join(ROOT_DIR, 'labels', 'val') if not os.path.exists(val_label): os.makedirs(val_label) # 生成测试集 test_image = os.path.join(ROOT_DIR, 'images', 'test') if not os.path.exists(test_image): os.makedirs(test_image) test_label = os.path.join(ROOT_DIR, 'labels', 'test') 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]
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。