#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Module for converting coco data set to detectron2 
   Usage:
   > python Coco2Detectron.py <inputJSONFile> <outputJSONFile>
   
   Input arguments:
   * <inputJSONFile> - a path to json file in coco format
   * <outputJSONFile> - name for the output file in detectron2 format
 
"""

__author__ = "Ondrej Klima"
__copyright__ = "Copyright 2020"
__credits__ = ["Ondrej Klima"]
__email__ = "iklima@fit.vutbr.cz"
__license__ = "BUT"
__version__ = "1.0"
__maintainer__ = "Ondrej Klima"

import json
import os
import sys
import os.path
from os import path

def main():
    argv = sys.argv 
    
    try:
        coco_filename = argv[1]
    except IndexError:
        raise IndexError('Input file in coco format must be supplied as an argument')
    if not path.isfile(coco_filename):
        raise ValueError('File "%s" does not exist!' % coco_filename)

    try:
        detectron_filename = argv[2]
    except IndexError:
        raise IndexError('Name for the output file must be supplied as an argument')
        
    with open(coco_filename, 'r') as coco_file:
        coco_data = json.load(coco_file) 

        dataset_dicts = []

        # Convert image list to dictionary 'image_id' => 'filename'
        for image in coco_data['images']:
            
            record = {}
            record["file_name"] = image['file_name']
            record["image_id"] = image['id']
            record["height"] = image['height']
            record["width"] = image['width']

            objects = []

            # Get annotations related to the image
            for annotation in coco_data['annotations']:
                if annotation['image_id'] == image['id']:
                
                    # convert from XYWH_ABS to XYXY_ABS
                    bbox = annotation['bbox']
                    for i in range(2):
                        bbox[i + 2] += bbox[i]

                    obj = {
                        "bbox": annotation['bbox'],
                        "bbox_mode": 0,
                        "segmentation": [annotation['segmentation'][0]],
                        "category_id": annotation['category_id'],
                        "iscrowd": 0
                        }
                    objects.append(obj)

            record["annotations"] = objects
            dataset_dicts.append(record)                  

        with open(detectron_filename, 'w') as output_json_file:
            json.dump(dataset_dicts, output_json_file)

if __name__ == "__main__":
    main()
