根据图片名生成COCO格式json文件（选择coco数据集的部分图像，生成其注释集）

先说自己写这个的动机，由于在利用coco数据集做一些工作，想做一些有针对性的任务，比如自己选择一些什么形状的目标对象或是目标对象为几个的工作，在网上也看到了根据图片名生成单个图像的json文件，还有下载部分COCO数据集并生成新的json标注文件，还有提取自己所需的类并生成新的json文件，其中第一个跟我们的需求是最像的，但是由于其提取的方式及内容与我们所要的有点不同，所以最后还是自己写了这个可以根据自己所选图像名称来提取其的json注释文件。

from pycocotools.coco import COCO
import os
import shutil
import json

'''
 路径参数
'''
#原coco数据集的路径
dataDir= '/root/biscuits/polar/data/coco/single-val'
#用于保存新生成的数据的路径
savepath = "data/"
#最后生产的json文件的保存路径
anno_save = savepath+'annotations/'
json_file='/root/biscuits/polar/data/coco/annotations/instances_val2017.json' # # json源文件

coco=COCO(json_file)
data=json.load(open(json_file,'r'))
'''
数据集参数
'''
#这里写自己需要提取的图像的文件夹位置
datasets_list=['train2017']
imagename = [f for f in os.listdir(os.path.join('/root/biscuits/polar/data/coco/single-val'))]
#生成保存路径
def mkr(path):
    if os.path.exists(path):
         shutil.rmtree(path)
         os.mkdir(path)
    else:
        os.mkdir(path)

#获取并处理所有需要的json数据
def process_json_data(annFile):
    imgIds_list = []
    anns_list = []
    print(imagename)
    for name_index in range(0, len(imagename)):
        # 通过imgID 找到其所有instance
        imgID = 0
        for i in range(0, len(data['images'])):
            if data['images'][i]['file_name'] == imagename[name_index]:  # 根据图片名找到对应的json中的'images'
                print(imagename[name_index])
                imgID1 = data['images'][i]['id']
                imgIds_list.append(imgID1)
                print(imgIds_list)
                for ann in data['annotations']:  # 根据image_id找到对应的annotation
                    if ann['image_id'] == imgID1:
                        anns_list.append(ann)

    image_info_list = coco.loadImgs(imgIds_list)
    classes_list = data['categories']
    return classes_list,image_info_list,anns_list

#保存数据到json
def save_json_data(json_file,classes_list,image_info_list,anns_list):
    coco_sub = dict()
    coco_sub['info'] = dict()
    coco_sub['licenses'] = []
    coco_sub['images'] = []
    coco_sub['type'] = 'instances'
    coco_sub['annotations'] = []
    coco_sub['categories'] = []
    #以下非必须,为coco数据集的前缀信息
    coco_sub['info']['description'] = 'COCO 2017 sub Dataset'
    coco_sub['info']['url'] = 'https://www.cnblogs.com/lhdb/'
    coco_sub['info']['version'] = '1.0'
    coco_sub['info']['year'] = 2020
    coco_sub['info']['contributor'] = 'smh'
    coco_sub['info']['date_created'] = '2020-7-1 10:06'
    sub_license = dict()
    sub_license['url'] =  'https://www.cnblogs.com/lhdb/'
    sub_license['id'] = 1
    sub_license['name'] = 'Attribution-NonCommercial-ShareAlike License'
    coco_sub['licenses'].append(sub_license)
    #以下为必须插入信息,包括image、annotations、categories三个字段
    #插入image信息
    coco_sub['images'].extend(image_info_list)
    #插入annotation信息
    coco_sub['annotations'].extend(anns_list)
    #插入categories信息
    coco_sub['categories'].extend(classes_list)
    #自此所有该插入的数据就已经插入完毕啦٩( ๑╹ ꇴ╹)۶
    #最后一步，保存数据
    json.dump(coco_sub, open(json_file, 'w'))


if __name__ == '__main__':
    mkr(anno_save)
    #按单个数据集进行处理
    for dataset in datasets_list:
        #获取要处理的json文件路径
        annFile='/root/biscuits/polar/data/coco/annotations/instances_val2017.json'.format(dataDir,dataset)
        #存储处理完成的json文件路径
        json_file = '/root/biscuits/polar/data/coco/annotations/mini_instances_val2017.json'.format(anno_save,dataset)
        #处理数据
        classes_list,image_info_list,anns_list = process_json_data(annFile)
        #保存数据
        save_json_data(json_file,classes_list,image_info_list,anns_list)
        print('instances_{}_sub.json'.format(dataset))

目前仍然小白，欢迎大家一起讨论。