import gradio as gr
import requests
from transformers import pipeline
from PIL import Image, ImageDraw
import numpy as np

obj_detector = pipeline(
    "object-detection", model="bninaos/finetuned-ViT-model",
device = 'cpu')

def plot_results(image, results, threshold=0.7):
    results = obj_detector(image)
    image = Image.fromarray(np.uint8(image))
    draw = ImageDraw.Draw(image)
    for result in results:
        score = result["score"]
        label = result["label"]
        box = list(result["box"].values())
        if score > threshold:
            x, y, x2, y2 = tuple(box)
            draw.rectangle((x, y, x2, y2), outline="red", width=1)
            draw.text((x, y), label, fill="white")
            draw.text(
                (x + 0.5, y - 0.5),
                text=str(score),
                fill="green" if score > threshold else "red",
            )
    return image

iface = gr.Interface(
    fn=plot_results,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Hardhat Detection",
    description="Upload an image to detect hardhats.",
)

iface.launch(share = True)