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

# Load the object detection pipeline using the DETR model
pipe = pipeline("object-detection", model="facebook/detr-resnet-50")

# Function to run object detection and return the image with bounding boxes
def detect_objects(image):
    # Get object detection results
    results = pipe(image)
    
    # Draw bounding boxes on the image
    draw = ImageDraw.Draw(image)
    
    for result in results:
        box = result['box']
        label = result['label']
        score = result['score']
        # Draw bounding box
        draw.rectangle(
            [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])], 
            outline="red", width=3
        )
        # Draw label and score
        draw.text((box['xmin'], box['ymin'] - 10), f"{label} ({score:.2f})", fill="red")
    
    return image

# Set up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Object Detection using DETR (facebook/detr-resnet-50)")

    # Input: Upload an image
    image_input = gr.Image(label="Upload Image", type="pil")

    # Output: Display the image with detected objects and bounding boxes
    output_image = gr.Image(label="Detected Objects")

    # Button to trigger object detection
    detect_button = gr.Button("Detect Objects")

    # Link button click with object detection function
    detect_button.click(fn=detect_objects, inputs=image_input, outputs=output_image)

# Launch the Gradio app
demo.launch()