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

pipe = pipeline("object-detection" , model = "hustvl/yolos-tiny")

def detect_objects(image):
    results = pipe(image)
    draw = ImageDraw.Draw(image)

    for r in results: 
        box = r['box']
        draw.rectangle(
            [(box['xmin'], box['ymin']) , (box['xmax'],box['ymax'])], outline = "blue" , width = 4 
        )

        draw.text((box['xmin'], box['ymin']), f"{r['label']} {r['score']:.2f}" , fill = "blue")

    for r in results: 
        print(f"{r['label']} ({r['score']:.2f}) - Box: {r['box']}")

    return image

with gr.Blocks(title = "🎯 Object Detection" , theme = gr.themes.Base()) as demo: 
    
    gr.Markdown("## 🎯Object Detection")
    
    with gr.Row(): 
        with gr.Column(): 
            image_input = gr.Image(label = "Upload an Image" , type = "pil")
            submit_btn = gr.Button("🔍 Detect Objects")
        with gr.Column(): 
            output_image = gr.Image(label= "Result" , type = "pil")

    submit_btn.click(
        fn = detect_objects , 
        inputs = [input_image] ,
        outputs= [output_image]
    )

demo.launch()