Create app.py

app.py

@@ -0,0 +1,56 @@
+import torch
+import gradio as gr
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
+import spaces
+
+# Use GPU if available
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
+processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+def query_image(Upload_Image, Text, score_threshold):
+    Text = Text
+    Text = Text.split(",")
+
+    size = max(Upload_Image.shape[:2])
+    target_sizes = torch.Tensor([[size, size]])
+    inputs = processor(text=Text, images=Upload_Image, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    outputs.logits = outputs.logits.cpu()
+    outputs.pred_boxes = outputs.pred_boxes.cpu()
+    results = 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
+        result_labels.append((box, Text[label.item()]))
+    return Upload_Image, result_labels
+
+
+description = """
+You can use AnyVision to query images with text descriptions of any object.
+To use it, simply upload an image and enter comma separated text descriptions of objects you want to query the image for. You
+can also use the score threshold slider to set a threshold to filter out low probability predictions.
+
+You can get better predictions by querying the image with text templates used in training the original model: e.g. *"photo of a star-spangled banner"*,
+*"image of a shoe"*.
+"""
+demo = gr.Interface(
+    query_image,
+    inputs=[gr.Image(), "text", gr.Slider(0, 1, value=0.1)],
+    outputs="annotatedimage",
+    title="AnyVision - Zero-Shot Object Detector with Owl2",
+    description=description
+)
+demo.launch()
+