说明
实现对目标检测数据集进行训练集和验证集的划分,自定义划分比例。图片可以包括jpg和png两种格式,标签是json格式。
目录结构
输入文件夹目录格式:
- data
1.1 images
--------------1.jpg
--------------2.jpg
--------------n.jpg
1.2 labels
--------------1.json
--------------2.json
--------------n.json
输出文件夹目录格式:
- 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)