按比例划分coco数据集,保存对应的json文件

一、pandas是什么?

经常要对数据集进行按比例划分,因为只有一个完整的coco数据文件,所以编写了这个脚本,可以按比例划分并保存成两个json文件。

二、完整代码

import os
import json
import shutil


def split_and_save_json(input_file, train_output_file, test_output_file, split_ratio=0.8):
    with open(input_file, 'r') as file:
        data = json.load(file)

    total_images = len(data['images'])
    split_index = int(total_images * split_ratio)

    train_data = {
        "info": data["info"],
        "images": data['images'][:split_index],
        "annotations": [ann for ann in data['annotations'] if ann['image_id'] <= data['images'][split_index - 1]['id']],
        "licenses": data["licenses"],
        "categories": data["categories"]
    }

    test_data = {
        "info": data["info"],
        "images": data['images'][split_index:],
        "annotations": [ann for ann in data['annotations'] if ann['image_id'] > data['images'][split_index - 1]['id']],
        "licenses": data["licenses"],
        "categories": data["categories"]
    }

    with open(train_output_file, 'w') as train_outfile:
        json.dump(train_data, train_outfile)

    with open(test_output_file, 'w') as test_outfile:
        json.dump(test_data, test_outfile)

    print(f"已将训练数据保存至 {train_output_file},并将测试数据保存至 {test_output_file}。")

# 调用新的函数,按照8:2划分数据并保存两个json文件
input_coco_file_path = '/codeyard/yolov5_6.0/data_MS_train_val/coco_train.json'
train_output_path = '/codeyard/yolov5_6.0/data_MS_train_val/coco_train.json'
test_output_path = '/codeyard/yolov5_6.0/data_MS_train_val/coco_test.json'
split_ratio = 0.8
split_and_save_json(input_coco_file_path, train_output_path, test_output_path, split_ratio)

首先,导入了 json 模块,该模块用于读取和写入 JSON 数据。

然后,定义了一个名为 split_and_save_json 的函数,它接受四个参数:input_file(输入文件路径)、train_output_file(训练集输出文件路径)、test_output_file(测试集输出文件路径)和可选的 split_ratio(分割比例,默认为 0.8)。

在函数内部,使用 open 函数打开输入文件,并使用 json.load 函数将文件内容加载为 Python 字典形式的数据。

获取图像总数和分割索引,分割索引是根据给定的分割比例计算得出的。

创建 train_data 字典,包含了训练集的信息,如数据集信息、前一部分图像、相应的注释、许可证和类别信息。

创建 test_data 字典,包含了测试集的信息,如数据集信息、后一部分图像、相应的注释、许可证和类别信息。

使用 open 函数分别打开训练集输出文件和测试集输出文件,并使用 json.dump 函数将相应的数据写入文件中。

最后,打印出保存的训练数据和测试数据的文件路径。

你可能感兴趣的:(目标检测,人工智能,机器学习,图像处理)