Application file

app.py

@@ -32,15 +32,15 @@ if torch.cuda.is_available():
 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")
-model_sam = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
-processor_sam = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+# model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
+# processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+# model_sam = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+# processor_sam = SamProcessor.from_pretrained("facebook/sam-vit-huge")
     
-# model = Owlv2ForObjectDetection.from_pretrained("owlv2-base-patch16-ensemble").to(device)
-# processor = Owlv2Processor.from_pretrained("owlv2-base-patch16-ensemble")
-# model_sam = SamModel.from_pretrained("SAM/sam-vit-huge").to(device)
-# processor_sam = SamProcessor.from_pretrained("SAM/sam-vit-huge")
+model = Owlv2ForObjectDetection.from_pretrained("owlv2-base-patch16-ensemble").to(device)
+processor = Owlv2Processor.from_pretrained("owlv2-base-patch16-ensemble")
+model_sam = SamModel.from_pretrained("SAM/sam-vit-huge").to(device)
+processor_sam = SamProcessor.from_pretrained("SAM/sam-vit-huge")
 
 
 @spaces.GPU
@@ -75,13 +75,13 @@ def query_image(img, text_queries, score_threshold=0.5):
     return sam_image,result_labels
 
    
-def generate_image_with_sam(img, boxes):
+def generate_image_with_sam(img, input_boxes):
     img_pil = Image.fromarray(img.astype('uint8'), 'RGB')
     inputs = processor_sam(img_pil, return_tensors="pt").to(device)
     
     image_embeddings = model_sam.get_image_embeddings(inputs["pixel_values"])
 
-    inputs = processor_sam(img_pil, input_boxes=[boxes], return_tensors="pt").to(device)
+    inputs = processor_sam(img_pil, input_boxes=[input_boxes], return_tensors="pt").to(device)
     inputs["input_boxes"].shape
     inputs.pop("pixel_values", None)
     inputs.update({"image_embeddings": image_embeddings})
@@ -101,16 +101,16 @@ Split anythings
 """
 demo = gr.Interface(
     fn=query_image,
-    inputs=[gr.Image(), gr.Textbox(label="Query Text"), gr.Slider(0, 1, value=0.5, label="Score Threshold")],
+    inputs=[gr.Image(), gr.Textbox(label="Query Text"), gr.Slider(0, 1, value=0.1, label="Score Threshold")],
     outputs=gr.AnnotatedImage(),
     title="Zero-Shot Object Detection SV3",
     description="This interface demonstrates object detection using  zero-shot object detection and SAM  for image segmentation.",
-    examples=[
-        ["images/purple cell.png", "purple cells", 0.05],
-        ["images/dark_cell.png", "gray cells", 0.1],
-        ["images/animals.png", "Rabbit,Squirrel,Parrot,Hedgehog,Turtle,Ladybug,Chick,Frog,Butterfly,Snail,Mouse", 0.35],
+    # examples=[
+    #     ["images/purple cell.png", "purple cells", 0.05],
+    #     ["images/dark_cell.png", "gray cells", 0.1],
+    #     ["images/animals.png", "Rabbit,Squirrel,Parrot,Hedgehog,Turtle,Ladybug,Chick,Frog,Butterfly,Snail,Mouse", 0.35],
 
-    ],
+    # ],
 )
 
 demo.launch()