add model

model_segvol_single.py

@@ -44,14 +44,16 @@ class SegVolModel(PreTrainedModel):
                 image, 
                 zoomed_image=None, 
                 text_prompt=None, 
-                bbox_prompt=None, 
-                bbox_prompt_map=None,
-                point_prompt=None, 
-                point_prompt_map=None,
+                bbox_prompt_group=None, 
+                point_prompt_group=None, 
                 use_zoom=True):
-        # print(image.shape, zoomed_image.shape, text_prompt)
         assert image.shape[0] == 1 and zoomed_image.shape[0] == 1, 'batch size should be 1'
-        # print(bbox_prompt.shape, bbox_prompt_map.shape, point_prompt[0].shape, point_prompt[1].shape, point_prompt_map.shape)
+        assert not (text_prompt is None and bbox_prompt_group is None and point_prompt_group is None)
+        bbox_prompt, bbox_prompt_map, point_prompt, point_prompt_map=None, None, None, None
+        if bbox_prompt_group is not None:
+            bbox_prompt, bbox_prompt_map = bbox_prompt_group
+        if point_prompt_group is not None:
+            point_prompt, point_prompt_map = point_prompt_group
         volume_shape = image[0][0].shape
 
         with torch.no_grad():