关于目标检测中按照比例将数据集随机划分成训练集和测试集

1. 前言

在做目标检测任务的时候,不少网上的数据,没有划分数据集,只是将数据和标签放在不同的文件夹下,没有划分数据集

虽然代码简单,每次重新编写还是颇为麻烦,这里记录一下

如下,有的数据集这样摆放:

关于目标检测中按照比例将数据集随机划分成训练集和测试集_第1张图片

这里的py文件是划分代码,和rawDataSet 放在相同目录下

关于图像分类任务划分数据集:关于图像分类任务中划分数据集,并且生成分类类别的josn字典文件

关于xml文件生成相应的类别json字典文件:

目标检测篇:如何根据xml标注文件生成类别classes的json文件

目标检测数据的可视化:

xml : 关于目标检测任务中,XML(voc格式)标注文件的可视化

txt : 关于目标检测任务中,YOLO(txt格式)标注文件的可视化

2. 完整代码

如下:

import random
import os
import shutil
from tqdm import tqdm


# 划分数据函数
def split_data(root,test_rate):
    images_path = [os.path.join(root,i) for i in os.listdir(root)]      # 获取所有图片路径
    test_split_path = random.sample(images_path, k=int(len(images_path) * test_rate))   # 随机采样测试集

    # 遍历所有图片
    for image_path in tqdm(images_path,desc='loading...'):
        # 获取相应的标注文件,这里需要根据目录、后缀更改
        label_path = image_path.replace('images','annotation')
        label_path = label_path.replace('.png','.xml')

        # 划分数据
        if image_path in test_split_path:       # 在测试集
            shutil.copy(image_path,'./data/test/images')
            shutil.copy(label_path,'./data/test/labels')
        else:
            shutil.copy(image_path,'./data/train/images')
            shutil.copy(label_path,'./data/train/labels')


if __name__ == '__main__':
    rawDataSet = './rawDataSet/images'            # 原始数据的图片路径

    if os.path.exists('./data'):        # 如果之前有,那么删除
        shutil.rmtree('./data')

    os.makedirs('./data/train/images')      # 训练集图片
    os.makedirs('./data/train/labels')      # 训练集标签
    os.makedirs('./data/test/images')       # 测试集图片
    os.makedirs('./data/test/labels')       # 测试集标签

    # 划分数据
    split_data(root=rawDataSet,test_rate=0.2)

代码运行过程:这里测试的就是五个数据

生成结果:

关于目标检测中按照比例将数据集随机划分成训练集和测试集_第2张图片

3. 代码使用的相关问题

不同于其他任务,脚本的运行逻辑是传入原始数据的图片目录,因为目标检测的标注文件和原图只是存放的目录不同,并且后缀不同,文件名是完全相同的

所以划分数据的时候,只需要找到图片,根据目录可以直接找到对应的标签

如下,第一个是目录的替换,根据数据进行更换,

第二个是图片后缀,根据自己数据集更换即可

如果是yolo标注的txt文件,也是一样的,后缀更换就行了

你可能感兴趣的:(#,关于,object,detection,目标检测,人工智能,机器学习)