合并多个coco格式数据的json标注文件

合并多个coco格式数据集的json标注文件

1. 拼接多个coco格式的json(单个coco json中可包含多个图片与标注)
2. 仅支持单个category
3. 会自动将 image id 与 bbox id 顺序排列
4. 仅支持 bbox，不支持 segmentaion

import json
import os
from tqdm import tqdm

if __name__ == '__main__':
    ''' 
        拼接多个coco格式的json(单个coco json中可包含多个图片与标注)
        仅支持单个category
        会自动将image id 与 bbox id 顺序排列
    '''
    path = 'coco_label'         # 改成coco label 的文件夹 
    output_file = 'test.json'   # 输出的文件名&路径
    categories = [                 #修改类别名
                     {
                         "supercategory": "animal",
                         "id": 1,
                         "name": "pig"
                     }
                 ]
    json_list = [i for i in os.listdir(path) if i.split('.')[-1] == 'json']
    json_list.sort()
    total_data = {}
    img_id_count, bbox_id_count = 0, 0

    total_data["images"], total_data["categories"], total_data["annotations"] = [], categories, []
    for js in json_list:
        print('正在处理json文件：', js)
        json_file = os.path.join(path, js)
        with open(json_file, 'r') as f:
            sub_data = json.load(f)
        for image in tqdm(sub_data['images']):
            if image['file_name'].count('/') == 4:
                image['file_name'] = image['file_name'].split('/')[-2] + '/' + image['file_name'].split('/')[-1]
            origin_img_id = image['id']
            image['id'] = img_id_count
            total_data["images"].append(image)
            for anno in sub_data['annotations']:
                if anno['image_id'] == origin_img_id:
                    anno['id'] = bbox_id_count
                    anno['image_id'] = img_id_count
                    total_data['annotations'].append(anno)
                    bbox_id_count += 1
            img_id_count += 1

    with open(output_file, 'w') as f:
        json.dump(total_data, f)