【目标检测】COCO数据集划分：训练集/验证集/测试集

原创已于 2023-05-18 12:48:25 修改 · 5.2k 阅读

43 ·

CC 4.0 BY-SA版权

文章标签：

#目标检测 #深度学习 #机器学习

于 2023-05-18 12:16:48 首次发布

目标检测专栏收录该内容

16 篇文章

订阅专栏

该文介绍了如何对COCO数据集进行训练、验证和测试集的划分，并使用Python代码示例展示了详细步骤。同时，文章提及了COCO数据集的目标检测标注可视化方法，帮助理解数据集的标注信息。

部署运行你感兴趣的模型镜像

1. COCO数据集划分

import os
import json
import numpy as np
import shutil
 
# 数据集路径
dataset_root = "E:\\MyDataset"
images_folder = os.path.join(dataset_root, "JPEGImages")
annotations_path = os.path.join(dataset_root, "annotations.json")
 
# 输出路径
output_root = os.path.join(dataset_root, "output")
os.makedirs(output_root, exist_ok=True)
 
# 读取annotations.json文件
with open(annotations_path, "r") as f:
    annotations_data = json.load(f)
 
# 提取images, annotations, categories
images = annotations_data["images"]
annotations = annotations_data["annotations"]
categories = annotations_data["categories"]
 
# 随机打乱数据
np.random.shuffle(images)
 
# 训练集，验证集，测试集比例
train_ratio, val_ratio, test_ratio = 0.7, 0.1, 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:]
 
# 分别为训练集、验证集和测试集创建子文件夹
train_folder = os.path.join(output_root, "train")
val_folder = os.path.join(output_root, "val")
test_folder = os.path.join(output_root, "test")
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
 
# 将图片文件复制到相应的子文件夹
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"]))
 
# 根据图片id分配annotations
def filter_annotations(annotations, image_ids):
    return [ann for ann in annotations if ann["image_id"] in image_ids]
 
train_ann = filter_annotations(annotations, [img["id"] for img in train_images])
val_ann = filter_annotations(annotations, [img["id"] for img in val_images])
test_ann = filter_annotations(annotations, [img["id"] for img in test_images])
 
# 生成train.json, val.json, test.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}
 
with open(os.path.join(output_root, "train.json"), "w") as f:
    json.dump(train_json, f)
 
with open(os.path.join(output_root, "val.json"), "w") as f:
    json.dump(val_json, f)
 
with open(os.path.join(output_root, "test.json"), "w") as f:
    json.dump(test_json, f)
 
print("数据集划分完成！")