分割COCO Format的目标识别数据集

最近需要利用coco format的数据集做目标识别的训练，个人比较熟悉tensorflow，但是tensorflow的Object Detection API接受的是tfrecord格式数据，为了分别生成train validation使用的tfrecord，我需要把coco format的数据集分割一下，在这里遇到了一些麻烦。

在网上找到了别人做好的轮子：

https://github.com/ashnair1/COCO-Assistant

但是问题在于，这个操作数据集的工具居然只能merge，不能split，无奈找了GitHub上这位同学的代码，对其中一些内容做了修改，最终实现了需求。

from pycocotools.coco import COCO
from pycocotools import mask as cocomask
import numpy as np
import skimage.io as io
import matplotlib.pyplot as plt
import pylab
import random
import os
from pycococreatortools import pycococreatortools as pt
import json
import random
root_dir = '/root'

INFO = {
    "description": "Construction Site Image Dataset",
    "url": "",
    "version": "0.2.0",
    "year": 2020,
    "contributor": "",
}

LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        "id": 1,
        "name": "truck",
        "supercategory": "class"
    },
    {
        "id": 2,
        "name": "excavator",
        "supercategory": "class"
    },
    {
        "id": 3,
        "name": "crane",
        "supercategory": "class"
    },
    {
        "id": 4,
        "name": "other_machine",
        "supercategory": "class"
    },
    {
        "id": 5,
        "name": "precast_concrete",
        "supercategory": "class"
    },
    {
        "id": 6,
        "name": "steel",
        "supercategory": "class"
    },
    {
        "id": 7,
        "name": "aggregates",
        "supercategory": "class"
    },
    {
        "id": 8,
        "name": "timber",
        "supercategory": "class"
    },
    {
        "id": 9,
        "name": "other_materials",
        "supercategory": "class"
    },
    {
        "id": 10,
        "name": "personnel",
        "supercategory": "class"
    },
]
coco_output_train = {
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}
coco_output_val = {
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}
coco_output_test = {
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}
annotation_root_path = '/root'
annotation_name = os.path.join(
    annotation_root_path, "via_export_coco_300_refine.json")
# dataset = json.load(open(annotation_name, 'r'))
#for ann in dataset['annotations']:
#    if not 'category_id' in ann.keys():
#        print(ann['id'])   # 改标签的时候总有遗漏，导致后面步骤出问题，这里check一下，如果没有输出数字说明没有遗漏
coco = COCO(annotation_name)
catogry_info = coco.loadCats(coco.getCatIds(
    catNms=["truck", "excavator", "crane","other_machine","precast_concrete","steel","aggregates","timber","other_materials","personnel"]))
coco_output_train['categories'] = catogry_info
coco_output_val['categories'] = catogry_info
coco_output_test['categories'] = catogry_info
image_ids = coco.getImgIds()[0:300]
random.shuffle(image_ids) #随机取样，最后再用

for image_id in image_ids[:200]:
    image_info = coco.loadImgs(image_id)[0]
    coco_output_train["images"].append(image_info)
    # annotation_ids = coco.getAnnIds(imgIds=image_id) 这句调用的函数里面有个迭代器有问题，不用这个，自己写一个
    annotation_ids = []
    for anno in coco.imgToAnns[str(image_id)]:
        annotation_ids.append(anno['id'])
    for annotation_id in annotation_ids:
        annotation_info = coco.loadAnns(annotation_id)[0]
        # 防止接下来的步骤再出错，输出的文件中image_id一律不再是字符串
        annotation_info['image_id'] = int(annotation_info['image_id'])
        coco_output_train["annotations"].append(annotation_info)

for image_id in image_ids[200:250]:
    image_info = coco.loadImgs(image_id)[0]
    coco_output_val["images"].append(image_info)
    # annotation_ids = coco.getAnnIds(imgIds=image_id) 这句调用的函数里面有个迭代器有问题，不用这个，自己写一个
    annotation_ids = []
    for anno in coco.imgToAnns[str(image_id)]:
        annotation_ids.append(anno['id'])
    for annotation_id in annotation_ids:
        annotation_info = coco.loadAnns(annotation_id)[0]
        # 防止接下来的步骤再出错，输出的文件中image_id一律不再是字符串
        annotation_info['image_id'] = int(annotation_info['image_id'])
        coco_output_val["annotations"].append(annotation_info)

for image_id in image_ids[250:]:
    image_info = coco.loadImgs(image_id)[0]
    coco_output_test["images"].append(image_info)
    # annotation_ids = coco.getAnnIds(imgIds=image_id) 这句调用的函数里面有个迭代器有问题，不用这个，自己写一个
    annotation_ids = []
    for anno in coco.imgToAnns[str(image_id)]:
        annotation_ids.append(anno['id'])
    for annotation_id in annotation_ids:
        annotation_info = coco.loadAnns(annotation_id)[0]
        # 防止接下来的步骤再出错，输出的文件中image_id一律不再是字符串
        annotation_info['image_id'] = int(annotation_info['image_id'])
        coco_output_test["annotations"].append(annotation_info)

with open('{}/anno_train_refine.json'.format(annotation_root_path), 'w') as output_json_file:
    json.dump(coco_output_train, output_json_file)
with open('{}/anno_val_refine.json'.format(annotation_root_path), 'w') as output_json_file:
    json.dump(coco_output_val, output_json_file)
with open('{}/anno_test_refine.json'.format(annotation_root_path), 'w') as output_json_file:
    json.dump(coco_output_test, output_json_file)