import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import AutoImageProcessor, AutoModelForObjectDetection

description = """
## This interface is made with 🤗 Gradio. 

Simply upload an image of any person wearning/not-wearing helmet. 
"""

image_processor = AutoImageProcessor.from_pretrained(
    "devonho/detr-resnet-50_finetuned_cppe5"
)
model = AutoModelForObjectDetection.from_pretrained(
    "devonho/detr-resnet-50_finetuned_cppe5"
)

# Gradio Components

image_in = gr.components.Image()
image_out = gr.components.Image()


def model_inference(img):
    with torch.no_grad():
        inputs = image_processor(images=img, return_tensors="pt")
        outputs = model(**inputs)
        target_sizes = torch.tensor([img.size[::-1]])
        results = image_processor.post_process_object_detection(
            outputs, threshold=0.5, target_sizes=target_sizes
        )[0]
    return results


def plot_results(image):
    image = Image.fromarray(np.uint8(image))
    results = model_inference(img=image)
    draw = ImageDraw.Draw(image)
    for score, label, box in zip(
        results["scores"], results["labels"], results["boxes"]
    ):
        score = score.item()
        box = [round(i, 2) for i in box.tolist()]
        x, y, x2, y2 = tuple(box)
        draw.rectangle((x, y, x2, y2), outline="red", width=1)
        draw.text((x, y), model.config.id2label[label.item()], fill="white")
        draw.text((x+0.5, y-0.5), text=str(score), fill='green' if score > 0.7 else 'red')
    return image


Iface = gr.Interface(
    fn=plot_results,
    inputs=[image_in],
    outputs=image_out,
    title="Object Detection Using Fine-Tuned Vision Transformers",
    description=description,
).launch()