add model

model_segvol_single.py

@@ -48,8 +48,10 @@ class SegVolModel(PreTrainedModel):
                 point_prompt=None, 
                 use_zoom=True,
                 **kwargs):
-        bbox_prompt, bbox_prompt_map = bbox_prompt
-        point_prompt, point_prompt_map = point_prompt
+        if bbox_prompt is not None:
+            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)
         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)
@@ -66,17 +68,51 @@ class SegVolModel(PreTrainedModel):
             size=volume_shape, mode='nearest')
         if not use_zoom:
             return logits_global_single
-        logits_global_single = logits_global_single[0][0]
         print(logits_global_single.shape)
         if point_prompt_map is not None:
             binary_points = F.interpolate(
-                point_prompt_map.unsqueeze(0).unsqueeze(0).float(),
-                size=volume_shape, mode='nearest')[0][0]
+                point_prompt_map.float(),
+                size=volume_shape, mode='nearest')
         if bbox_prompt_map is not None:
             binary_cube = F.interpolate(
-                bbox_prompt_map.unsqueeze(0).unsqueeze(0).float(),
-                size=volume_shape, mode='nearest')[0][0]
-        return
+                bbox_prompt_map.float(),
+                size=volume_shape, mode='nearest')
+        
+        min_d, min_h, min_w, max_d, max_h, max_w = logits2roi_coor(self.config.spatial_size, logits_global_single[0][0])
+        assert min_d is None, 'Fail to detect foreground!'
+
+        # Crop roi
+        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)
+        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]
+            prompt_reflection = (
+                binary_cube_cropped,
+                global_preds
+            )
+        if point_prompt is not None:
+            binary_points_cropped = binary_points[:, :, min_d:max_d+1, min_h:max_h+1, min_w:max_w+1]
+            prompt_reflection = (
+                binary_points_cropped,
+                global_preds
+            )
+    
+        ## inference
+        with torch.no_grad():
+            logits_single_cropped = sliding_window_inference(
+                    image_single_cropped.cuda(), prompt_reflection,
+                    self.config.spatial_size, 1, self.model, 0.5,
+                    text=text_prompt,
+                    use_box=bbox_prompt is not None,
+                    use_point=point_prompt is not None,
+                )
+            logits_single_cropped = logits_single_cropped.cpu().squeeze()
+        logits_global_single[:, :, min_d:max_d+1, min_h:max_h+1, min_w:max_w+1] = logits_single_cropped
+        return logits_global_single
+
 
 
 # processor