updating the detection model from owl-vit to owl-v2

app.py

@@ -83,10 +83,10 @@ def run_image_captioner(image, device):
 
 @spaces.GPU()
 def run_segmentation(image, object_to_segment, device):
-    # OWL-ViT for object detection
-    owl_vit_model_id = "google/owlvit-base-patch32"
-    processor = OwlViTProcessor.from_pretrained(owl_vit_model_id)
-    od_model = OwlViTForObjectDetection.from_pretrained(owl_vit_model_id).to(device)
+    # OWL-V2 for object detection
+    owl_v2_model_id = "google/owlv2-base-patch16-ensemble"
+    processor = Owlv2Processor.from_pretrained(owl_vit_model_id)
+    od_model = Owlv2ForObjectDetection.from_pretrained(owl_vit_model_id).to(device)
     text_queries = [object_to_segment]
     inputs = processor(text=text_queries, images=image, return_tensors="pt").to(device)
     with torch.no_grad():