add model

model_segvol_single.py

@@ -40,29 +40,44 @@ class SegVolModel(PreTrainedModel):
 
         self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-    def forward(self, image, zoomed_image=None, text_prompt=None, bbox_prompt=None, point_prompt=None, **kwargs):
+    def forward_test(self, 
+                image, 
+                zoomed_image=None, 
+                text_prompt=None, 
+                bbox_prompt=None, 
+                point_prompt=None, 
+                use_zoom=True,
+                **kwargs):
         bbox_prompt, bbox_prompt_map = bbox_prompt
         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)
         volume_shape = image[0][0].shape
-        # test mode
-        if self.config.test_mode:
-            with torch.no_grad():
-                logits_global_single = self.model(zoomed_image,
-                                                    text=text_prompt, 
-                                                    boxes=bbox_prompt, 
-                                                    points=point_prompt)
-            print(logits_global_single.shape)
-            logits_global_single = F.interpolate(
-                logits_global_single.cpu(),
+
+        with torch.no_grad():
+            logits_global_single = self.model(zoomed_image,
+                                                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
+        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]
-            print(logits_global_single.shape)
-            return
-        else:
-            print('unsupport training mode now')
-            return
-        return self.model.forward(image, text=None, boxes=None, points=None, **kwargs)
+        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
+
 
 # processor
 class SegVolProcessor():
@@ -136,6 +151,25 @@ class SegVolProcessor():
         binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=label_single_resize.shape).unsqueeze(0).unsqueeze(0)
         return box_single, binary_cube_resize
 
+    def dice_score(self, preds, labels):
+        assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match\n" + str(preds.shape) + str(labels.shape)
+        predict = preds.view(1, -1)
+        target = labels.view(1, -1)
+        if target.shape[1] < 1e8:
+            predict = predict.to(self.custom_device)
+            target = target.to(self.custom_device)
+        predict = torch.sigmoid(predict)
+        predict = torch.where(predict > 0.5, 1., 0.)
+        
+        tp = torch.sum(torch.mul(predict, target))
+        den = torch.sum(predict) + torch.sum(target) + 1
+        dice = 2 * tp / den
+
+        if target.shape[1] < 1e8:
+            predict = predict.cpu()
+            target = target.cpu()
+        return dice
+    
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):
         d = dict(data)