import os
import argparse
import numpy as np
import cv2
import json
from tqdm import tqdm
import pycocotools.mask as mask_util

parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, required=True)
parser.add_argument('--output', type=str, required=True)
args = parser.parse_args()


def rle_to_polygon(rle) -> (list, list, float):
    mask = mask_util.decode(rle)
    mask = np.where(mask > 0.5, 1, 0)
    mask = np.ascontiguousarray(mask, dtype=np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
    polygon = []
    segs = []
    contours = [contour for contour in contours if contour.size >= 6]
    for contour in contours:
        _contour = contour.flatten().tolist()
        _contour.append(_contour[0])
        _contour.append(_contour[1])
        polygon.extend(_contour)
        segs.append(contour.flatten().tolist())
    RLEs = mask_util.frPyObjects(segs, mask.shape[0], mask.shape[1])
    RLE = mask_util.merge(RLEs)
    area = float(mask_util.area(RLE))
    return polygon, area


def main(args):
    with open(args.input, 'r') as f:
        ann_data = json.load(f)
    for ann in tqdm(ann_data['annotations'], total=len(ann_data['annotations'])):
        rle_data = ann['segmentation']
        # encodedRLE = mask_util.frPyObjects(rle_data, rle_data['size'][0], rle_data['size'][1])
        # polygon, area = rle_to_polygon(encodedRLE)
        polygon, area = rle_to_polygon(rle_data)
        ann['segmentation'] = [polygon]
        ann['area'] = area
        ann['iscrowd'] = 0
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    with open(args.output, 'w') as f:
        json.dump(ann_data, f)


if __name__ == '__main__':
    main(args)