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

"""Module for converting binary masks to vector representation in coco format
   Usage:
   > python Masks2Coco.py <inputImages> <productBinaryMasks> <packageBinaryMasks>
   <rangeFrom> <rangeTo> <outputJSONFile>
   
   Input arguments:
   * <inputImages> - a path to a directory containing input images
   * <productBinaryMasks> - directory containing input product binary masks
   * <packageBinaryMasks> - directory containing input package binary masks
   * <rangeFrom> - from which image in the data set to start
   * <rangeTo> - at which image at the data set to stop
   * <outputJSONFile> - output file in the Coco format
   
"""

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

import datetime
import json
import os, os.path
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools
import sys

INFO = {
    "description": "Tipp-Ex Corrector Dataset",
    "url": "https://www.fit.vutbr.cz/",
    "version": "1",
    "year": 2020,
    "contributor": "iklima",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "BUT Licence",
        "url": "https://www.fit.vutbr.cz/"
    }
]

CATEGORIES = [
    {
        'id': 0,
        'name': 'product',
        'supercategory': 'shape',
    },
    {
        'id': 1,
        'name': 'package',
        'supercategory': 'shape',
    },
]

def main():
    # Parsing input arguments 
    argv = sys.argv    

    try:
        inputImagesPath = argv[1]
    except IndexError:
        raise IndexError('Input directory path containing images must be supplied as an argument')

    try:
        productMasksPath = argv[2]
    except IndexError:
        raise IndexError('Input directory path containing product binary masks must be supplied as an argument')
        
    try:
        packageMasksPath = argv[3]
    except IndexError:
        raise IndexError('Input directory path containing package binary masks must be supplied as an argument')
        
    try:
        rangeFrom = argv[4]
    except IndexError:
        raise IndexError('A number of image in the data set from which to start')
        
    try:
        rangeTo = argv[5]
    except IndexError:
        raise IndexError('A number of image in the data at which to stop')

    try:
        outputFile = argv[6]
    except IndexError:
        raise IndexError('Output json file in the coco format')        

    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

    segmentation_id = 0
    image_id = 0

    for i in range(int(rangeFrom), int(rangeTo)):
        annotation1_filename = productMasksPath + '/%d.png' % i
        annotation2_filename = packageMasksPath + '/%d.png' % i
        image_filename = inputImagesPath + '/%d.png' % i

        image = Image.open(image_filename)
        image_info = pycococreatortools.create_image_info(
            image_id, os.path.basename(image_filename), image.size)
        coco_output["images"].append(image_info)

        category_info = {'id': 0, 'is_crowd': False}
        binary_mask = np.asarray(Image.open(annotation1_filename)
            .convert('1')).astype(np.uint8)
        
        annotation_info = pycococreatortools.create_annotation_info(
            segmentation_id, image_id, category_info, binary_mask,
            image.size, tolerance=2)

        if annotation_info is not None:
            coco_output["annotations"].append(annotation_info)           

        segmentation_id += 1

        category_info = {'id': 1, 'is_crowd': False}
        binary_mask = np.asarray(Image.open(annotation2_filename)
            .convert('1')).astype(np.uint8)
        
        annotation_info = pycococreatortools.create_annotation_info(
            segmentation_id, image_id, category_info, binary_mask,
            image.size, tolerance=2)

        if annotation_info is not None:
            coco_output["annotations"].append(annotation_info)           

        segmentation_id += 1

        image_id += 1
        print(image_id)
       

    with open(outputFile, 'w') as output_json_file:
        json.dump(coco_output, output_json_file)

if __name__ == "__main__":
    main()
