import gradio as gr
from PIL import ImageDraw
from transformers import pipeline

object_detector = None

def get_detector():
    global object_detector
    if object_detector is None:
        object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
    return object_detector

def draw_bounding_boxes(image, object_detections):
    draw = ImageDraw.Draw(image)
    for detection in object_detections:
        box = detection["box"]
        label = detection["label"]
        score = detection["score"]
        draw.rectangle(
            (box["xmin"], box["ymin"], box["xmax"], box["ymax"]),
            outline=(255, 0, 0),
            width=2,
        )
        draw.text((box["xmin"], box["ymin"] - 20), f"{label} ({score:.2f})", fill=(255, 0, 0))
    return image

def detect_object(image):
    output = get_detector()(image)
    return draw_bounding_boxes(image, output)

demo = gr.Interface(
    fn=detect_object,
    inputs=[gr.Image(label="Select Image", type="pil")],
    outputs=[gr.Image(label="Processed Image", type="pil")],
    title="@IT AI Enthusiast - Project 6: Object Detector",
    description="Detect objects in the provided input image.",
)
demo.launch()