Transform Tsinghua-Tencent 100K Dataset (ver 2021) Annotations to COCO Format

tt100k2021_to_coco.py

import json
import argparse
import copy


def load_json(file_name):
    file = open(file_name, 'r').read()
    return json.loads(file)


def parse(data):
    # File Format
    result_train = {
        "info": {
            "description": "TT100K Dataset (ver. 2021) COCO Format",
            "url": "https://github.com/zhaoweizhong",
            "version": "2.0",
            "year": 2021,
            "contributor": "Zhaowei Zhong",
            "date_created": "2021/03/05"
        },
        "licenses": [
            {
                "url": "https://creativecommons.org/licenses/by-nc-sa/4.0/",
                "id": 1,
                "name": "Attribution-NonCommercial-ShareAlike 4.0 International License (CC BY-NC-SA 4.0)"
            }
        ],
        "images": [],
        "annotations": [],
        "categories": []
    }

    # Categories
    categories = ["pl80", "w9", "p6", "ph4.2", "i8", "w14", "w33", "pa13", "im", "w58", "pl90", "il70", "p5", "pm55", "pl60", "ip", "p11", "pdd", "wc", "i2r", "w30", "pmr", "p23", "pl15", "pm10", "pss", "w1", "p4", "w38", "w50", "w34", "pw3.5", "iz", "w39", "w11", "p1n", "pr70", "pd", "pnl", "pg", "ph5.3", "w66", "il80", "pb", "pbm", "pm5", "w24", "w67", "w49", "pm40", "ph4", "w45", "i4", "w37", "ph2.6", "pl70", "ph5.5", "i14", "i11", "p7", "p29", "pne", "pr60", "pm13", "ph4.5", "p12", "p3", "w40", "pl5", "w13", "pr10", "p14", "i4l", "pr30", "pw4.2", "w16", "p17", "ph3", "i9", "w15", "w35", "pa8", "pt", "pr45", "w17", "pl30", "pcs", "pctl", "pr50", "ph4.4", "pm46", "pm35", "i15", "pa12", "pclr", "i1", "pcd", "pbp", "pcr", "w28", "ps", "pm8", "w18", "w2", "w52", "ph2.9", "ph1.8", "pe", "p20", "w36", "p10", "pn", "pa14", "w54", "ph3.2", "p2", "ph2.5", "w62", "w55", "pw3", "pw4.5", "i12", "ph4.3", "phclr", "i10", "pr5", "i13", "w10", "p26", "w26", "p8", "w5", "w42", "il50", "p13", "pr40", "p25", "w41", "pl20", "ph4.8", "pnlc", "ph3.3", "w29", "ph2.1", "w53", "pm30", "p24", "p21", "pl40", "w27", "pmb", "pc", "i6", "pr20", "p18", "ph3.8", "pm50", "pm25", "i2", "w22", "w47", "w56", "pl120", "ph2.8", "i7", "w12", "pm1.5", "pm2.5", "w32", "pm15", "ph5", "w19", "pw3.2", "pw2.5", "pl10", "il60", "w57", "w48", "w60", "pl100", "pr80", "p16", "pl110", "w59", "w64", "w20", "ph2", "p9", "il100", "w31", "w65", "ph2.4", "pr100", "p19", "ph3.5", "pa10", "pcl", "pl35", "p15", "w7", "pa6", "phcs", "w43", "p28", "w6", "w3", "w25", "pl25", "il110", "p1", "w46", "pn-2", "w51", "w44", "w63", "w23", "pm20", "w8", "pmblr", "w4", "i5", "il90", "w21", "p27", "pl50", "pl65", "w61", "ph2.2", "pm2", "i3", "pa18", "pw4"]
    i = 1
    for category in categories:
        result_train['categories'].append({
            "id": i,
            "name": category
        })
        i = i + 1

    result_test = copy.deepcopy(result_train)

    # Images and Annotations
    count = len(data['imgs'])
    count_train = int(count * 0.8)
    count_test = count - count_train
    
    i = 1
    anno_id = 1
    for img in data['imgs']:
        if i <= count_train:
            flag = False
            for box in data['imgs'][img]['objects']:
                if box['category'] in categories:
                    flag = True
            if flag:
                result_train['images'].append({
                    "license": 1,
                    "file_name": data['imgs'][img]['path'].split('/')[1],
                    "height": 2048,
                    "width": 2048,
                    "id": data['imgs'][img]['id']
                })
                for box in data['imgs'][img]['objects']:
                    if box['category'] in categories:
                        result_train['annotations'].append({
                            "segmentation": [[]],
                            "area": (box['bbox']['xmax'] - box['bbox']['xmin']) * (box['bbox']['ymax'] - box['bbox']['ymin']),
                            "iscrowd": 0,
                            "image_id": data['imgs'][img]['id'],
                            "bbox": [
                                box['bbox']['xmin'],
                                box['bbox']['ymin'],
                                box['bbox']['xmax'] - box['bbox']['xmin'],
                                box['bbox']['ymax'] - box['bbox']['ymin']
                            ],
                            "category_id": categories.index(box['category']) + 1,
                            "id": anno_id
                        })
                        anno_id += 1
                        if ('ellipse_org' in box):
                            for xy in box['ellipse_org']:
                                result_train['annotations'][i]['segmentation'][0].append(xy[0])
                                result_train['annotations'][i]['segmentation'][0].append(xy[1])
                        elif 'polygon' in box:
                            for xy in box['polygon']:
                                result_train['annotations'][i]['segmentation'][0].append(xy[0])
                                result_train['annotations'][i]['segmentation'][0].append(xy[1])
        else:
            flag = False
            for box in data['imgs'][img]['objects']:
                if box['category'] in categories:
                    flag = True
            if flag:
                result_test['images'].append({
                    "license": 1,
                    "file_name": data['imgs'][img]['path'].split('/')[1],
                    "height": 2048,
                    "width": 2048,
                    "id": data['imgs'][img]['id']
                })
                for box in data['imgs'][img]['objects']:
                    if box['category'] in categories:
                        result_test['annotations'].append({
                            "segmentation": [[]],
                            "area": (box['bbox']['xmax'] - box['bbox']['xmin']) * (box['bbox']['ymax'] - box['bbox']['ymin']),
                            "iscrowd": 0,
                            "image_id": data['imgs'][img]['id'],
                            "bbox": [
                                box['bbox']['xmin'],
                                box['bbox']['ymin'],
                                box['bbox']['xmax'] - box['bbox']['xmin'],
                                box['bbox']['ymax'] - box['bbox']['ymin']
                            ],
                            "category_id": categories.index(box['category']) + 1,
                            "id": anno_id
                        })
                        anno_id += 1
                        if ('ellipse_org' in box):
                            for xy in box['ellipse_org']:
                                result_test['annotations'][i]['segmentation'][0].append(xy[0])
                                result_test['annotations'][i]['segmentation'][0].append(xy[1])
                        elif 'polygon' in box:
                            for xy in box['polygon']:
                                result_test['annotations'][i]['segmentation'][0].append(xy[0])
                                result_test['annotations'][i]['segmentation'][0].append(xy[1])
        i = i + 1
    
    print('Train Images: ' + str(len(result_train['images'])))
    print('Test Images: ' + str(len(result_test['images'])))
    print('Train Annotations: ' + str(len(result_train['annotations'])))
    print('Test Annotations: ' + str(len(result_test['annotations'])))

    with open('train.json', "w") as f:
        json.dump(result_train, f)

    with open('test.json', "w") as f:
        json.dump(result_test, f)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--file_name', type=str, default='annotations_all.json')

    args = parser.parse_args()
    file_name = args.file_name

    data = load_json(file_name)
    parse(data)

you must add then 66 so 67 line and tab all codes inline
if data['imgs'][img]['path'].split('/')[0] == "train":

same way add 97. line
if data['imgs'][img]['path'].split('/')[0] == "test":