convert_annotation_to_coco_style.py

# this is a script to convert the pseudo label generated by export_pseudo_label_CC3M.py to COCO style dataset file

import json, argparse
from collections import Counter
parser = argparse.ArgumentParser("converter")
parser.add_argument('--annotation_file', type=str)
parser.add_argument('--output_file', type=str)
args = parser.parse_args()

with open(args.annotation_file) as f:
    annotation = json.load(f)

# fix the bug of multiple annotation for one instance
if isinstance(annotation['annotations'][0]['category'], list):
    current_index = 0
    current_image = -1
    for i, anno in enumerate(annotation['annotations']):
        if anno['image_id'] != current_image:
            current_image = anno['image_id']
            current_index = 0
        anno['category'] = anno['category'][current_index]
        current_index += 1

# 

# categories
# # create the field categories
categories = set()
inst_counter = Counter()
for instance in annotation['annotations']:
    categories.add(instance['category'])
    inst_counter[instance['category']] += 1

image_counter = Counter()
for image in annotation['images']:
    for k in image['pos_category_ids']:
        image_counter[k] += 1

cats = list()
cat_id = dict()
for id, cat in enumerate(categories):
    cat_id[cat] = id + 1
    cats.append(dict(
        id=id + 1,
        name=cat,
        instance_count=inst_counter[cat],
        image_count=image_counter[cat],
    ))
annotation['categories'] = cats

# images
# # remove the images that does not have an instance
annotation['images'] = [k for k in annotation['images'] if len(k['pos_category_ids']) > 0]
# # convert category name to category id
for image in annotation['images']:
    image['pos_category_ids'] = [cat_id[k] for k in image['pos_category_ids']]

# annotations
# category -> category id
for anno in annotation['annotations']:
    anno['category_id'] = cat_id[anno.pop('category')]

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