# =============================================================================
# Created By  : Kanaad Pathak
# Created Date: Fri 17 Jul 2020 07∶55∶25 PM BST 2020
# Course : MSc AIA 2020 | University of Strathclyde | United Kingdom
# =============================================================================

'''
Purpose : This module is used for drawing the bounding boxes for the JSON files generated by supervise.ly inference

'''
# =============================================================================


# =============================================================================
# IMPORT BLOCK
# =============================================================================

import json
import os
import cv2
import pandas as pd
from tqdm import tqdm
import numpy as np

# =============================================================================


'''
Function : Get Bounding Box
Input : JSON file containing annotations
Output : Bounding box coordinates for that file

Corner Cases : Only JSON file format accepted from supervisely. Other JSON inference formats may not work.

'''

def get_bbox(annotation_file):
    with open(annotation_file) as json_file:
        data = json.load(json_file)
        bbox = []
        num_objects = len(data['objects'])
        for i in range(num_objects):
            box = data['objects'][i]['points']['exterior']
            box = np.asarray(box)
            box = box.flatten()
            bbox.append(box)
    return bbox


'''
Function : Draw Rectangle
Input : Bounding box coordinates , thickness of line [integer value] , Concatenation of original image to generated image [boolean flag]
Output : New JPEG file containing drawn bounding boxes

Corner Cases : Image must correspond to bounding box locations, otherwise function errors out.

'''

def draw_rect(bbox,image,thickness,concat_orig):
    image = cv2.imread(image)
    original = image.copy()
    for i in range(len(bbox)):
        x1,x2,y1,y2 = bbox[i]
        image = cv2.rectangle(image,(x1,x2),(y1,y2),(0,0,255), thickness)
    if concat_orig:
        img_1 = original
        img_2 = image
        out_img = np.concatenate((img_1,img_2),axis=1)
        return out_img
    return image
'''
Function : Main
Input : None
Output : Folder containing all images with bounding boxes drawn on them

Corner Cases : All images and corresponding annotation files must be present in current directory. This can be manually done and hence no code was written.


'''

def main():
    annotations = os.listdir('ann')
    imgs = os.listdir('img')
    annotations.sort()
    imgs.sort()
    img_df = pd.DataFrame(list(zip(annotations, imgs)), columns =['annotation', 'image'])
    os.mkdir('outputs')

    for img in tqdm(range(len(img_df))):
        annotation_file = img_df['annotation'].iloc[img]
        image = img_df['image'].iloc[img]
        bbox = get_bbox(annotation_file)
        drawn_img = draw_rect(bbox,image,5,True)
        path = 'outputs/' + str(image)
        cv2.imwrite(path,drawn_img)

if __name__ == "__main__":
    main()

