import gradio as gr
from PIL import Image
from transformers import pipeline
from PIL import ImageDraw 
#..
#test
object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")

def draw_bounding_boxes(image, detections):
    draw = ImageDraw.Draw(image)
    
    for detection in detections:
        box = detection['box']
        label = detection['label']
        xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']
        
        draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)
        draw.text((xmin, ymin - 10), label, fill="red")
    
    return image

def detect_object(image):
    output = object_detector(image)  # Use 'image' instead of 'raw_image'
    processed_image = draw_bounding_boxes(image.copy(), output)  # Make a copy to avoid modifying the original
    return processed_image

gr.close_all()

demo = gr.Interface(
    fn=detect_object, 
    inputs=[gr.Image(label="Select Image", type="pil")], 
    outputs=[gr.Image(label="Image with Bounding Box", type="pil")], 
    title="Object Detector", 
    theme="soft",
    description="This is an object detection model that detects objects in an image and draws bounding boxes around them. "
)

demo.launch(share=True)