将COCO格式的物体检测数据集划分训练集、验证集和测试集

目录

导入所需库

定义数据集路径

创建输出目录

读取JSON注释文件

随机打乱图像列表

计算划分大小

复制图像到相应文件夹

完整代码

---

导入所需库

我们需要以下Python库：

os：处理文件路径。

json：读取和写入JSON文件。

numpy：随机打乱图像列表。

shutil：复制图像文件。

import os
import json
import numpy as np
import shutil

定义数据集路径

设置数据集的根目录、图像文件夹和注释文件路径。

根目录："D:\\dataset"

图像文件夹："D:\\dataset\\images"

注释文件："D:\\dataset\\annotations.json"

# 数据集路径（请根据实际情况修改）
dataset_root = "D:\\dataset"
images_folder = os.path.join(dataset_root, "images")
annotations_path = os.path.join(dataset_root, "annotations.json")

 

创建输出目录

在根目录下创建output文件夹，并在其中创建out_train、out_val和out_test子文件夹。

# 输出路径
output_root = os.path.join(dataset_root, "output")
os.makedirs(output_root, exist_ok=True)

train_folder = os.path.join(output_root, "out_train")
val_folder = os.path.join(output_root, "out_val")
test_folder = os.path.join(output_root, "out_test")
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)

读取JSON注释文件

加载COCO格式的JSON文件，提取images（图像信息）、annotations（标注信息）和categories（类别信息）。

# 读取注释文件
with open(annotations_path, "r") as f:
    annotations_data = json.load(f)

# 提取数据
images = annotations_data["images"]
annotations = annotations_data["annotations"]
categories = annotations_data["categories"]

随机打乱图像列表

使用numpy随机打乱图像列表，确保划分的随机性。

# 随机打乱图像列表
np.random.shuffle(images)

计算划分大小

根据图像总数和比例计算训练集和测试集的大小：

假设图像总数为N。

训练集：N * 0.8。

验证集：N * 0.0 = 0。

测试集：N * 0.2。

# 定义划分比例
train_ratio, val_ratio, test_ratio = 0.8, 0, 0.2

# 计算大小
num_images = len(images)
num_train = int(num_images * train_ratio)
num_val = int(num_images * val_ratio)  # 将为0

# 划分图像
train_images = images[:num_train]
val_images = images[num_train:num_train + num_val]  # 空列表
test_images = images[num_train + num_val:]

复制图像到相应文件夹

将训练集和测试集的图像复制到对应的文件夹。

# 复制图像
for img in train_images:
    shutil.copy(os.path.join(images_folder, img["file_name"]), 
                os.path.join(train_folder, img["file_name"]))

for img in val_images:  # 不会执行
    shutil.copy(os.path.join(images_folder, img["file_name"]), 
                os.path.join(val_folder, img["file_name"]))

for img in test_images:
    shutil.copy(os.path.join(images_folder, img["file_name"]), 
                os.path.join(test_folder, img["file_name"]))

完整代码

以下是完整的Python脚本：

import os
import json
import numpy as np
import shutil# 数据集路径（请根据实际情况修改）
dataset_root = "D:\\dataset"
images_folder = os.path.join(dataset_root, "images")
annotations_path = os.path.join(dataset_root, "annotations.json")# 输出路径
output_root = os.path.join(dataset_root, "output")
os.makedirs(output_root, exist_ok=True)train_folder = os.path.join(output_root, "out_train")
val_folder = os.path.join(output_root, "out_val")
test_folder = os.path.join(output_root, "out_test")
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)# 读取注释文件
with open(annotations_path, "r") as f:annotations_data = json.load(f)# 提取数据
images = annotations_data["images"]
annotations = annotations_data["annotations"]
categories = annotations_data["categories"]# 随机打乱图像列表
np.random.shuffle(images)# 定义划分比例
train_ratio, val_ratio, test_ratio = 0.8, 0, 0.2# 计算大小
num_images = len(images)
num_train = int(num_images * train_ratio)
num_val = int(num_images * val_ratio)# 划分图像
train_images = images[:num_train]
val_images = images[num_train:num_train + num_val]
test_images = images[num_train + num_val:]# 复制图像
for img in train_images:shutil.copy(os.path.join(images_folder, img["file_name"]), os.path.join(train_folder, img["file_name"]))for img in val_images:shutil.copy(os.path.join(images_folder, img["file_name"]), os.path.join(val_folder, img["file_name"]))for img in test_images:shutil.copy(os.path.join(images_folder, img["file_name"]), os.path.join(test_folder, img["file_name"]))# 函数：过滤注释
def filter_annotations(annotations, image_ids):return [ann for ann in annotations if ann["image_id"] in image_ids]# 获取image_ids
train_image_ids = [img["id"] for img in train_images]
val_image_ids = [img["id"] for img in val_images]
test_image_ids = [img["id"] for img in test_images]# 过滤注释
train_ann = filter_annotations(annotations, train_image_ids)
val_ann = filter_annotations(annotations, val_image_ids)
test_ann = filter_annotations(annotations, test_image_ids)# 创建JSON字典
train_json = {"images": train_images, "annotations": train_ann, "categories": categories}
val_json = {"images": val_images, "annotations": val_ann, "categories": categories}
test_json = {"images": test_images, "annotations": test_ann, "categories": categories}# 写入JSON文件
with open(os.path.join(output_root, "out_train.json"), "w") as f:json.dump(train_json, f)
with open(os.path.join(output_root, "out_val.json"), "w") as f:json.dump(val_json, f)
with open(os.path.join(output_root, "out_test.json"), "w") as f:json.dump(test_json, f)print("数据集划分完成！")