# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import json
import colorsys
import pandas as pd, numpy as np

from matplotlib import font_manager
from PIL import Image, ImageDraw, ImageFont

from project_utils import (
    get_color_map,
    convert_box,
    load_image,
    save_image
)

df = dataiku.Dataset("test_scored").get_dataframe()
folder = dataiku.Folder("aKlMxTsk")
output_folder = dataiku.Folder("LkxFClYj")

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
labels = set()
for i in range(len(df)):
    label = json.loads(df.iloc[i].label)
    for d in label:
        labels.add(d["category"])

color_map = get_color_map(list(labels))

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
font = font_manager.FontProperties()
file = font_manager.findfont(font)
font = ImageFont.truetype(file, 10)

for idx in range(len(df)):

    image = load_image(folder, df.iloc[idx].record_id)
    boxes = json.loads(df.iloc[idx].label)
    pred_boxes = json.loads(df.iloc[idx].prediction)

    dctx = ImageDraw.Draw(image)

    # Add predicted bounding boxes
    for i in range(len(pred_boxes)):
        pred_box, label = convert_box(pred_boxes[i]["bbox"]), pred_boxes[i]["category"]
        color = color_map[label]
        dctx.rectangle(list(pred_box), outline=color)
        text_size = font.getbbox(label)[2:]
        dctx.rectangle((*pred_box[:2], pred_box[0]+text_size[0]+2, pred_box[1]+text_size[1]+2), fill=color)
        dctx.text((pred_box[0]+1, pred_box[1]+1), label, fill="white", font=font)

    # Add ground truth bounding boxes
    for i in range(len(boxes)):
        color = color_map[boxes[i]["category"]]
        dctx.rectangle(convert_box(boxes[i]["bbox"]), outline=color)

    save_image(output_folder, df.iloc[idx].record_id, image)