数据集生成 YOLOV5 可训练的数据目录、并且可视化

1、前言

YOLOV5 训练数据的目录结构如下：

如果有测试集的话，也按照下面目录摆放即可

注意：这里的图片和标签文件名要严格对应！！后缀除外

关于YOLOv5介绍或者yolo格式的介绍参考之前专栏，

2、划分数据+生成YOLOV5的训练目录+可视化

测试的数据目录结构如下：

2.1 划分数据保存成yolov5的目录

代码如下：

import os
import shutil
import random
from tqdm import tqdm'''
--datasets--images--train       # 训练集图片--images--val         # 验证集图片--datasets--labels--train       # 训练集标签--labels--val         # 验证集标签
'''def mkdir():os.mkdir('datasets')os.mkdir('datasets/images')os.mkdir('datasets/labels')os.mkdir('datasets/images/train')os.mkdir('datasets/images/val')os.mkdir('datasets/labels/train')os.mkdir('datasets/labels/val')def split_yolov5(floder,rate,img_f):mkdir()         # 创建 yolov5的目录，已有datasets文件夹，需要删除floder = os.path.join(floder,'images')image_path = [os.path.join(floder,i) for i in os.listdir(floder)]val_path = random.sample(image_path,k=int(len(image_path)*rate))        # 划分的 val数据train_num = 0val_num = 0for i in tqdm(image_path):label_path = i.replace(('.'+img_f),'.txt')      # 标签路径label_path = label_path.replace('images','labels')if i in val_path:       # 验证集val_num +=1shutil.copy(i,'./datasets/images/val')shutil.copy(label_path,'./datasets/labels/val')else:                   # 训练集train_num +=1shutil.copy(i,'./datasets/images/train')shutil.copy(label_path,'./datasets/labels/train')print('split over!!')print('data set number is: ',len(image_path))print('train set number is: ',train_num)print('val set number is: ',val_num)if __name__ == '__main__':root = 'data'       # 待划分的数据目录split_rate = 0.2    # 划分验证集的比例image_format = 'jpg'        # 数据图片的格式split_yolov5(floder=root,rate=split_rate,img_f=image_format)

参数按照注释填即可

控制台输出：

生成的数据：

2.2 可视化

根据之前的yolo可视化参考代码内容：关于目标检测任务中，YOLO(txt格式)标注文件的可视化_yolo格式的标签-CSDN博客

这里只是将目录更改：

需要提供txt的classes类别文本！！

import cv2
import os
import randomdef txtShow(img, txt, save=True):image = cv2.imread(img)height, width = image.shape[:2]  # 获取原始图像的高和宽# 读取classes类别信息with open('datasets/classes.txt', 'r') as f:classes = f.read().splitlines()# ['Leconte', 'Boerner', 'linnaeus', 'armandi', 'coleoptera', 'acuminatus', 'Linnaeus']# 读取yolo格式标注的txt信息with open(txt, 'r') as f:labels = f.read().splitlines()# ['0 0.403646 0.485491 0.103423 0.110863', '1 0.658482 0.425595 0.09375 0.099702', '2 0.482515 0.603795 0.061756 0.045387', '3 0.594122 0.610863 0.063244 0.052083', '4 0.496652 0.387649 0.064732 0.049107']ob = []  # 存放目标信息for i in labels:cl, x_centre, y_centre, w, h = i.split(' ')# 需要将数据类型转换成数字型cl, x_centre, y_centre, w, h = int(cl), float(x_centre), float(y_centre), float(w), float(h)name = classes[cl]  # 根据classes文件获取真实目标xmin = int(x_centre * width - w * width / 2)  # 坐标转换ymin = int(y_centre * height - h * height / 2)xmax = int(x_centre * width + w * width / 2)ymax = int(y_centre * height + h * height / 2)tmp = [name, xmin, ymin, xmax, ymax]  # 单个检测框ob.append(tmp)# 绘制检测框for name, x1, y1, x2, y2 in ob:cv2.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)  # 绘制矩形框cv2.putText(image, name, (x1, y1 - 10), fontFace=cv2.FONT_HERSHEY_SIMPLEX,fontScale=0.5, thickness=1, color=(0, 0, 255))# 保存图像if save:cv2.imwrite('result.png', image)# 展示图像cv2.imshow('test', image)cv2.waitKey()cv2.destroyAllWindows()if __name__ == '__main__':img_path = './datasets/images/train/'  # 传入图片image = [os.path.join(img_path, i) for i in os.listdir(img_path)]r = random.randint(0, len(image) - 1)  # 生成随机索引image_path = image[r]labels_path = image_path.replace('images', 'labels')  # 自动获取对应的 txt 标注文件labels_path = labels_path.replace('.jpg', '.txt')txtShow(img=image_path, txt=labels_path, save=True)

需要更改的就是图片后缀：

展示：

3、生成YOLOV5的训练目录

对于已经划分好的数据集，直接copy即可

待拷贝的目录：因为之前本人习惯用下面的目录结构，所以这里也是一样

完整代码：

import os
import shutil
import random
from tqdm import tqdm'''
--datasets--images--train       # 训练集图片--images--val         # 验证集图片
--datasets--labels--train       # 训练集标签--labels--val         # 验证集标签
'''def mkdir():os.mkdir('datasets')os.mkdir('datasets/images')os.mkdir('datasets/labels')os.mkdir('datasets/images/train')os.mkdir('datasets/images/val')os.mkdir('datasets/labels/train')os.mkdir('datasets/labels/val')def split_yolov5(floder,img_f):mkdir()  # 创建 yolov5的目录，已有datasets文件夹，需要删除train_floder = os.path.join(floder, 'train/images')image_path = [os.path.join(train_floder, i) for i in os.listdir(train_floder)]train_num = 0for i in tqdm(image_path,desc='train set copy'):label_path = i.replace(('.' + img_f), '.txt')  # 标签路径label_path = label_path.replace('images', 'labels')train_num += 1shutil.copy(i, './datasets/images/train')shutil.copy(label_path, './datasets/labels/train')val_floder = os.path.join(floder, 'val/images')val_path = [os.path.join(val_floder, i) for i in os.listdir(val_floder)]val_num = 0for i in tqdm(val_path,desc='val set copy'):label_path = i.replace(('.' + img_f), '.txt')  # 标签路径label_path = label_path.replace('images', 'labels')val_num += 1shutil.copy(i, './datasets/images/val')shutil.copy(label_path, './datasets/labels/val')print('copy over!!')print('train set number is: ', train_num)print('val set number is: ', val_num)if __name__ == '__main__':root = 'data'  # 划分好的数据目录image_format = 'jpg'  # 数据图片的格式split_yolov5(floder=root,img_f=image_format)

控制台输出：