import cv2
import khandy
import numpy as np
import gradio as gr
from PIL import Image
from insectid import InsectDetector
from insectid import InsectIdentifier


def inference(filename):
    detector = InsectDetector()
    identifier = InsectIdentifier()
    image = khandy.imread(filename)

    if image is None:
        return None

    if max(image.shape[:2]) > 1280:
        image = khandy.resize_image_long(image, 1280)

    image_for_draw = image.copy()
    image_height, image_width = image.shape[:2]
    boxes, confs, classes = detector.detect(image)

    for box, _, _ in zip(boxes, confs, classes):
        box = box.astype(np.int32)
        box_width = box[2] - box[0] + 1
        box_height = box[3] - box[1] + 1

        if box_width < 30 or box_height < 30:
            continue

        cropped = khandy.crop_or_pad(image, box[0], box[1], box[2], box[3])
        results = identifier.identify(cropped)
        print(results[0])
        prob = results[0]['probability']

        if prob < 0.10:
            text = 'Unknown'
        else:
            text = '{} {}: {:.2f}%'.format(
                results[0]['chinese_name'],
                results[0]['latin_name'],
                100.0 * results[0]['probability']
            )

        position = [box[0] + 2, box[1] - 20]
        position[0] = min(max(position[0], 0), image_width)
        position[1] = min(max(position[1], 0), image_height)
        cv2.rectangle(
            image_for_draw,
            (box[0], box[1]),
            (box[2], box[3]),
            (0, 255, 0),
            2
        )

        image_for_draw = khandy.draw_text(
            image_for_draw,
            text,
            position,
            font='simsun.ttc',
            font_size=15
        )

    return Image.fromarray(image_for_draw[:, :, ::-1], mode='RGB')


with gr.Blocks() as demo:
    with gr.Tab("Image"):
        gr.Markdown("## Insect Inference on Image")
        with gr.Row():
            image_input = gr.Image(
                type='filepath',
                label="Input Image",
                source="upload"
            )

            image_output = gr.Image(
                type='pil',
                label="Output Image",
                source="canvas"
            )

        text_button = gr.Button("Detect")

    text_button.click(inference, inputs=image_input, outputs=image_output)


demo.launch()