app.py

import torch
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
from PIL import Image, ImageDraw
import gradio as gr

checkpoint = "google/owlvit-base-patch32"

model = AutoModelForZeroShotObjectDetection.from_pretrained(checkpoint)
processor = AutoProcessor.from_pretrained(checkpoint)

def detect_objects(image, text_queries):

    if isinstance(image, str):
        image = Image.open(image)

    inputs = processor(images=image, text=text_queries, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        target_sizes = torch.tensor([image.size[::-1]])
        results = processor.post_process_object_detection(outputs, threshold=0.1, target_sizes=target_sizes)[0]

    draw = ImageDraw.Draw(image)
    
    scores = results["scores"].tolist()
    labels = results["labels"].tolist()
    boxes = results["boxes"].tolist()

    for box, score, label in zip(boxes, scores, labels):
        xmin, ymin, xmax, ymax = box
        draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=1)
        draw.text((xmin, ymin), f"{text_queries[label]}: {round(score, 2)}", fill="black")

    return image

inputs = [
    gr.Image(type="pil", label="Input Image"),
    gr.Textbox(label="Text Queries (comma-separated)")
]

output = gr.Image(type="pil", label="Output Image")

gr.Interface(
    fn=detect_objects,
    inputs=inputs,
    outputs=output,
    title="Zero-Shot Object Detection",
    description="Detect objects in an image using zero-shot object detection.",
).launch()