Create app.py

app.py

@@ -0,0 +1,110 @@
+import torch
+import gradio as gr
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
+import spaces
+
+# --------------------------
+# Device setup
+# --------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --------------------------
+# Load OWLv2 model
+# --------------------------
+model = Owlv2ForObjectDetection.from_pretrained(
+    "google/owlv2-base-patch16-ensemble"
+).to(device)
+
+processor = Owlv2Processor.from_pretrained(
+    "google/owlv2-base-patch16-ensemble"
+)
+
+# --------------------------
+# Detection function
+# --------------------------
+@spaces.GPU
+def query_image(img, text_queries, score_threshold):
+
+    # Convert query string to list
+    text_queries = [q.strip() for q in text_queries.split(",")]
+
+    # ✅ FIX: Use actual image size
+    h, w = img.shape[:2]
+    target_sizes = torch.tensor([[h, w]])
+
+    # Prepare inputs
+    inputs = processor(
+        text=text_queries,
+        images=img,
+        return_tensors="pt"
+    ).to(device)
+
+    # Run model
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # Move outputs to CPU
+    outputs.logits = outputs.logits.cpu()
+    outputs.pred_boxes = outputs.pred_boxes.cpu()
+
+    # Post-process predictions
+    results = processor.post_process_object_detection(
+        outputs=outputs,
+        target_sizes=target_sizes
+    )
+
+    boxes = results[0]["boxes"]
+    scores = results[0]["scores"]
+    labels = results[0]["labels"]
+
+    detections = []
+
+    for box, score, label in zip(boxes, scores, labels):
+
+        if score < score_threshold:
+            continue
+
+        x1, y1, x2, y2 = box.tolist()
+
+        detections.append({
+            "box": [round(x1, 2), round(y1, 2), round(x2, 2), round(y2, 2)],
+            "label": text_queries[label.item()],
+            "score": round(float(score), 3)
+        })
+
+    return img, detections
+
+
+# --------------------------
+# Gradio UI
+# --------------------------
+demo = gr.Interface(
+    fn=query_image,
+    inputs=[
+        gr.Image(type="numpy", label="Upload Image"),
+        gr.Textbox(
+            label="Enter objects (comma separated)",
+            value="person, car, dog"
+        ),
+        gr.Slider(
+            minimum=0,
+            maximum=1,
+            value=0.2,
+            step=0.01,
+            label="Score Threshold"
+        )
+    ],
+    outputs=gr.AnnotatedImage(label="Detection Results"),
+    title="OWLv2 Zero-Shot Object Detection",
+    description=(
+        "Upload an image and type objects to detect.\n\n"
+        "Example: 'person, car, dog'\n\n"
+        "Tip: Use natural phrases like 'photo of a car' for better results."
+    )
+)
+
+# --------------------------
+# Run app
+# --------------------------
+if __name__ == "__main__":
+    demo.launch()