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

news2024/10/6 0:30:57

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 = 0
    val_num = 0
    for 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 +=1
            shutil.copy(i,'./datasets/images/val')
            shutil.copy(label_path,'./datasets/labels/val')

        else:                   # 训练集
            train_num +=1
            shutil.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 random


def 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 = 0
    for 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 += 1
        shutil.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 = 0
    for 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 += 1
        shutil.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)

控制台输出：