Create app.py

app.py

@@ -0,0 +1,78 @@
+import spaces
+from transformers import Owlv2Processor, Owlv2ForObjectDetection, AutoProcessor, AutoModelForZeroShotObjectDetection
+import torch
+import gradio as gr
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+owl_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
+owl_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+dino_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
+dino_model = AutoModelForZeroShotObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base").to(device)
+
+@spaces.GPU
+def infer(img, text_queries, score_threshold, model):
+  
+  if model == "dino":
+    queries=""
+    for query in text_queries:
+      queries += f"{query}. "
+
+    width, height = img.shape[:2]
+
+    target_sizes=[(width, height)]
+    inputs = dino_processor(text=queries, images=img, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+      outputs = dino_model(**inputs)
+      outputs.logits = outputs.logits.cpu()
+      outputs.pred_boxes = outputs.pred_boxes.cpu()
+      results = dino_processor.post_process_grounded_object_detection(outputs=outputs, input_ids=inputs.input_ids,
+                                                                    box_threshold=score_threshold,
+                                                                    target_sizes=target_sizes)
+  elif model == "owl":
+    size = max(img.shape[:2])
+    target_sizes = torch.Tensor([[size, size]])
+    inputs = owl_processor(text=text_queries, images=img, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+      outputs = owl_model(**inputs)
+      outputs.logits = outputs.logits.cpu()
+      outputs.pred_boxes = outputs.pred_boxes.cpu()
+      results = owl_processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
+      
+  boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
+  result_labels = []
+
+  for box, score, label in zip(boxes, scores, labels):
+      box = [int(i) for i in box.tolist()]
+      if score < score_threshold:
+          continue
+      if model == "owl":
+        label = text_queries[label.cpu().item()]
+      result_labels.append((box, label))
+  return result_labels
+
+def query_image(img, text_queries, owl_threshold, dino_threshold):
+    text_queries = text_queries
+    text_queries = text_queries.split(",")
+    owl_output = infer(img, text_queries, owl_threshold, "owl")
+    dino_output = infer(img, text_queries, owl_threshold, "dino")
+
+
+    return (img, owl_output), (img, dino_output)
+
+
+owl_threshold = gr.Slider(0, 1, value=0.16, label="OWL Threshold")
+dino_threshold = gr.Slider(0, 1, value=0.12, label="Grounding DINO Threshold")
+owl_output = gr.AnnotatedImage(label="OWL Output")
+dino_output = gr.AnnotatedImage(label="Grounding DINO Output")
+demo = gr.Interface(
+    query_image,
+    inputs=[gr.Image(label="Input Image"), gr.Textbox("Candidate Labels"), owl_threshold, dino_threshold],
+    outputs=[owl_output, dino_output],
+    title="Zero-Shot Object Detection with OWLv2",
+    examples=[["./bee.jpg", "bee, flower", 0.16, 0.12]]
+)
+demo.launch(debug=True)