'''
Author: your name
Date: 2021-10-03 13:47:29
LastEditTime: 2021-10-11 17:36:56
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: \cxf\merge_multipart_coco.py
'''
# !pyinstaller -F merge_multipart_coco.py -i merge_multipart_coco.ico
# chenxinfeng
# ------使用方法------
# 直接拖动json到EXE中
# merge parts of 'rat1', 'rat2' into 'rat's
#


# %% import
import json
import sys, os
import datetime
from pycocotools.coco import COCO

def convert(anno, anno_out):
    coco = COCO(anno)
    dataout = dict()
    dataout['categories'] = coco.dataset['categories']
    dataout['images'] = []
    dataout['annotations'] = []
    dataout['info'] = dict(
        description = 'labelme2posecoco', 
        version = '1.0', 
        year = datetime.datetime.today().year,
        date_created = datetime.datetime.today().strftime('%Y/%m/%d'))
    dataout['licences'] = ''
    dataout['images'] = coco.dataset['images']
    # %%
    annoid =0
    imgids = coco.getImgIds()
    catids = list(coco.cats.keys())
    for image_id in imgids:
        anns = coco.imgToAnns[image_id]
        allrat_seg = {catid:[] for catid in catids}

        for iann in anns:
            catid = iann['category_id']
            allrat_seg[catid] += iann['segmentation']

        for catid, segs in allrat_seg.items():
            if segs == []: continue
            segs_1d = [j for sub in segs for j in sub]
            x_min, y_min = min(segs_1d[::2]), min(segs_1d[1::2])
            x_max, y_max = max(segs_1d[::2]), max(segs_1d[1::2])
            bbox = [int(i) for i in [x_min, y_min, x_max-x_min, y_max-y_min]]
            annotations0 = dict(
                id = annoid,
                bbox = bbox,
                iscrowd = 0,
                area = bbox[2]*bbox[3],
                category_id = catid,
                image_id = image_id,
                segmentation = segs
            )
            dataout['annotations'].append(annotations0)
            annoid+=1

    with open(anno_out, 'w') as f:
        json.dump(dataout, f, indent=4)

if __name__ == "__main__":
    assert len(sys.argv)==2
    anno = sys.argv[1]
    assert os.path.isfile(anno)
    prefix =  os.path.splitext(anno)[0]
    anno_out = prefix + '_merge.json'
    print(anno)
    convert(anno, anno_out)