按比例划分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 函数将相应的数据写入文件中。

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