# clean up the merged COCO dataset
from os.path import join as pjoin
import json
import pycocotools.mask as mask_util

ROOT = 'data/coco_merged/annotations'
JSONS = ['interaction.json', 'interactable.json', 'semantics.json']

# filter noise produced by labeling tool, denote mask in compressed RLE format
def fix_mask(data):
    for anno in data['annotations']:
        mask = mask_util.frPyObjects(anno['segmentation'], 540, 960)
        mask = mask_util.decode(mask)
        mask[:, -1] = mask[:, -2]
        mask[-1:, :] = mask[-2, :]
        RLE = mask_util.encode(mask)
        bbox = mask_util.toBbox(RLE)
        bbox = bbox.astype(int)
        area = int(bbox[2] * bbox[3])
        RLE['counts'] = RLE['counts'].decode('utf-8')
        anno['segmentation'] = RLE
        anno['bbox'] = bbox.tolist()
        anno['area'] = area

# remove redundant annotations
def remove_redundant(data):
    total = len(data['annotations'])
    anno_list = []

    for anno in data['annotations']:
        anno['id']=0
        if anno not in anno_list:
            anno_list.append(anno)
    
    for anno in anno_list:
        anno['id'] = anno_list.index(anno) + 1

    data['annotations'] = anno_list
    print(f'{total - len(anno_list)} redundant annotations removed')


def main():
    json_files = [pjoin(ROOT, j) for j in JSONS]
    for j in json_files:
        with open(j, 'r') as f:
            data = json.load(f)
            fix_mask(data)
            remove_redundant(data)
        new_j = j.replace('.json', '_cleaned.json')
        with open(new_j, 'w') as f:
            json.dump(data, f, indent=4)


if __name__ == '__main__':
    main()