COCOAPI

# This is a sample Python script.

# from pycocotools.coco import COCO
import json

from shapely.geometry import Polygon
from shapely.ops import cascaded_union

import os
import re
import argparse


def get_children(path=os.curdir):
    result = []
    for parent, _, files in os.walk(path):
        for name in files:
            # result.append(os.path.join(os.path.abspath(parent), name))
            result.append(os.path.join(parent, name))
    return result


def get_children_with_match(path=os.curdir, regexp='.*'):
    result = []
    for cursor in get_children(path):
        if not re.match(regexp, cursor) is None:
            result.append(cursor)
    return result


def points_to_polygon(points):
    return Polygon([(points[i], points[i+1]) for i in range(0, len(points), 2)])


def segmentation_to_polygons(segmentation):
    return [points_to_polygon(points) for points in segmentation]


def polygon_to_points(polygon):
    points = []
    for p in list(polygon.exterior.coords):
        points.append(p[0])
        points.append(p[1])
    return points


def union_polygons(input_json, output_json):
    print(f'Input json path: {input_json}')
    print(f'Output json path: {output_json}')

    with open(input_json, 'r') as file:
        content = file.read()
        coco = json.loads(content)

        polygons = []
        for ann in coco['annotations']:
            polygons += segmentation_to_polygons(ann['segmentation'])
        union_polygon = cascaded_union(polygons)

        coco['annotations'] = [{
            'area': union_polygon.area,
            'bbox': union_polygon.bounds,
            'category_id': 1,
            'iscrowd': 0,
            'segmentation': [polygon_to_points(union_polygon)]
        }]

        with open(output_json, 'x') as output_file:
            output_file.write(json.dumps(coco))


def main(input_dir, output_dir):
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)

    # `for i in $(find . -iname '*.json'); do iconv -f euc-kr -t utf-8 $i > ${i%.*}.utf8.json; done`
    files = get_children_with_match(input_dir, regexp=r'.*\.json')
    for f in files:
        names = os.path.splitext(os.path.basename(f))
        union_polygons(f, os.path.join(output_dir, names[0]+'.output'+names[1]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='RealTimeVideo demo')
    parser.add_argument(
        '--input',
        '-i',
        default='input',
        help='Input directory')
    parser.add_argument(
        '--output',
        '-o',
        default='output',
        help='Output directory')
    args = parser.parse_args()
    main(args.input, args.output)