目标检测数据集划分训练集和验证集

说明

实现对目标检测数据集进行训练集和验证集的划分，自定义划分比例。图片可以包括jpg和png两种格式，标签是json格式。

目录结构

输入文件夹目录格式：

1. data
   1.1 images
   --------------1.jpg
   --------------2.jpg
   --------------n.jpg
   1.2 labels
   --------------1.json
   --------------2.json
   --------------n.json

输出文件夹目录格式：

1. dataset
   —1.1 train
   -------1.1.1 images
   -------1.1.2 labels
   —1.2 val
   -------1.2.1 images
   -------1.2.2 labels

实现代码

import os
import json
import random
import shutil

# 设置随机种子，保证每次划分结果一致
random.seed(42)

# 指定原始数据集所在目录和划分后数据集存放目录
data_dir = "/home/znsz/wanfu/PaddleDetection/dataset/grape_730"
train_dir = "/home/znsz/wanfu/PaddleDetection/dataset/grape_dataset/train/"
val_dir = "/home/znsz/wanfu/PaddleDetection/dataset/grape_dataset/val/"

# 构建训练集和验证集目录
if not os.path.exists(train_dir):
    os.makedirs(train_dir)
if not os.path.exists(val_dir):
    os.makedirs(val_dir)

# 遍历数据集中的所有图像文件
for filename in os.listdir(data_dir + "/images"):
    if filename.endswith(".jpg") or filename.endswith(".png"):
        
        # 根据设定的比例随机划分到训练集或验证集
        if random.random() <= 0.8:
            dst_dir = train_dir
        else:
            dst_dir = val_dir
        
        # 构建图片文件路径和标签文件路径
        img_path = os.path.join(data_dir, "images", filename)
        label_path = os.path.join(data_dir, "labels", filename.replace(".jpg", ".json").replace(".png", ".json"))
        
        # 构建目标保存路径
        dst_img_path = os.path.join(dst_dir, "images", filename)
        dst_label_path = os.path.join(dst_dir, "labels", filename.replace(".jpg", ".json").replace(".png", ".json"))
        
        # 复制图片文件和标签文件到目标路径
        shutil.copy(img_path, dst_img_path)
        shutil.copy(label_path, dst_label_path)