add model

model_segvol_single.py

@@ -52,9 +52,9 @@ class SegVolModel(PreTrainedModel):
             bbox_prompt, bbox_prompt_map = bbox_prompt
         if point_prompt is not None:
             point_prompt, point_prompt_map = point_prompt
-        print(image.shape, zoomed_image.shape, text_prompt)
+        # 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)
+        # print(bbox_prompt.shape, bbox_prompt_map.shape, point_prompt[0].shape, point_prompt[1].shape, point_prompt_map.shape)
         volume_shape = image[0][0].shape
 
         with torch.no_grad():
@@ -62,14 +62,12 @@ class SegVolModel(PreTrainedModel):
                                                 text=text_prompt, 
                                                 boxes=bbox_prompt, 
                                                 points=point_prompt)
-        print(logits_global_single.shape)
         logits_global_single = F.interpolate(
             logits_global_single.cpu(),
             size=volume_shape, mode='nearest')
         if not use_zoom:
             return logits_global_single
-        print(logits_global_single.shape)
-        print(torch.unique(logits_global_single))
+        
         if point_prompt_map is not None:
             binary_points = F.interpolate(
                 point_prompt_map.float(),
@@ -88,7 +86,7 @@ class SegVolModel(PreTrainedModel):
         image_single_cropped = image[:, :, min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
         global_preds = (torch.sigmoid(logits_global_single[:, :, min_d:max_d+1, min_h:max_h+1, min_w:max_w+1])>0.5).long()
         
-        assert not (bbox_prompt is not None and point_prompt is not None)
+        assert not (bbox_prompt is not None and point_prompt is not None), 'Do not use point prompt and box prompt at the same time.'
         prompt_reflection = None
         if bbox_prompt is not None:
             binary_cube_cropped = binary_cube[:, :, min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]