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

objDetector = pipeline(task = "object-detection", 
                       model = "facebook/detr-resnet-50")

def detect_objects(image):           # image is a NumPy array
    image = Image.fromarray(image)   # convert from NumPy array to PIL image
    draw = ImageDraw.Draw(image)     # create a drawable image
    results = objDetector(image)     # detect objects in image
    
    for obj in results:
        if obj['score'] < 0.9:
            continue
        box = [
            obj['box']['xmin'],
            obj['box']['ymin'],
            obj['box']['xmax'],
            obj['box']['ymax']
        ]
        draw.rectangle(box, outline='yellow', width=5)
        draw.text((box[0], box[1] - 15), obj['label'], fill='white')
    
    return image

demo = gr.Interface(
    fn = detect_objects,
    inputs = gr.Image(label = "Upload Image"),
    outputs = gr.Image(label = "Detected Objects")
)

demo.launch()