add model

model_segvol_single.py

@@ -1,5 +1,5 @@
 from transformers import PreTrainedModel
-from .config_segvol import SegVolConfig
+from config_segvol import SegVolConfig
 import numpy as np
 import monai.transforms as transforms
 
@@ -77,17 +77,111 @@ class SegVolProcessor():
 
         # transform
         item = self.transform(item)
-        print('ready for zoom out')
         item_zoom_out = self.zoom_out_transform(item)
         item['zoom_out_image'] = item_zoom_out['image']
         item['zoom_out_label'] = item_zoom_out['label']
-        print(  'Zoom_in image shape: ',   item['image'].shape, 
-                '\nZoom_in label shape: ', item['label'].shape,
-                '\nZoom_out image shape: ', item['zoom_out_image'].shape,
-                '\nZoom_out label shape: ', item['zoom_out_label'].shape,
-            )
+        del item['image_transforms']
+        del item['label_transforms']
         return item
 
+
+    def generate_box(self, pred_pre, bbox_shift=None):
+        meaning_post_label = pred_pre # [h, w, d]
+        ones_idx = (meaning_post_label > 0).nonzero(as_tuple=True)
+        if all(tensor.nelement() == 0 for tensor in ones_idx):
+            bboxes = torch.tensor([-1,-1,-1,-1,-1,-1])
+            # print(bboxes, bboxes.shape)
+            return bboxes
+        min_coords = [dim.min() for dim in ones_idx]    # [x_min, y_min, z_min]
+        max_coords = [dim.max() for dim in ones_idx]    # [x_max, y_max, z_max]
+
+
+        if bbox_shift is None:
+            corner_min = []
+            corner_max = []
+            shape = meaning_post_label.shape
+            for coor in min_coords:
+                coor_ = max(0, coor)
+                corner_min.append(coor_)
+            for idx, coor in enumerate(max_coords):
+                coor_ = min(shape[idx], coor)
+                corner_max.append(coor_)
+            corner_min = torch.tensor(corner_min)
+            corner_max = torch.tensor(corner_max)
+            return torch.cat((corner_min, corner_max), dim=0)
+        else:
+            # add perturbation to bounding box coordinates
+            corner_min = []
+            corner_max = []
+            shape = meaning_post_label.shape
+            for coor in min_coords:
+                coor_ = max(0, coor + random.randint(-bbox_shift, bbox_shift))
+                corner_min.append(coor_)
+            for idx, coor in enumerate(max_coords):
+                coor_ = min(shape[idx], coor + random.randint(-bbox_shift, bbox_shift))
+                corner_max.append(coor_)
+            corner_min = torch.tensor(corner_min)
+            corner_max = torch.tensor(corner_max)
+            return torch.cat((corner_min, corner_max), dim=0)
+
+
+    def select_points(self, preds, num_positive_extra=4, num_negative_extra=0, fix_extra_point_num=None):
+        spacial_dim = 3
+        points = torch.zeros((0, 3))
+        labels = torch.zeros((0))
+        pos_thred = 0.9
+        neg_thred = 0.1
+        
+        # get pos/net indices
+        positive_indices = torch.nonzero(preds > pos_thred, as_tuple=True) # ([pos x], [pos y], [pos z])
+        negative_indices = torch.nonzero(preds < neg_thred, as_tuple=True)
+
+        ones_idx = (preds > pos_thred).nonzero(as_tuple=True)
+        if all(tmp.nelement() == 0 for tmp in ones_idx):
+            # all neg
+            num_positive_extra = 0
+            selected_positive_point = torch.tensor([-1,-1,-1]).unsqueeze(dim=0)
+            points = torch.cat((points, selected_positive_point), dim=0)
+            labels = torch.cat((labels, torch.tensor([-1]).reshape(1)))
+        else:
+            # random select a pos point
+            random_idx = torch.randint(len(positive_indices[0]), (1,))
+            selected_positive_point = torch.tensor([positive_indices[i][random_idx] for i in range(spacial_dim)]).unsqueeze(dim=0)
+            points = torch.cat((points, selected_positive_point), dim=0)
+            labels = torch.cat((labels, torch.ones((1))))
+
+        if num_positive_extra > 0:
+            pos_idx_list = torch.randperm(len(positive_indices[0]))[:num_positive_extra]
+            extra_positive_points = []
+            for pos_idx in pos_idx_list:
+                extra_positive_points.append([positive_indices[i][pos_idx] for i in range(spacial_dim)])
+            extra_positive_points = torch.tensor(extra_positive_points).reshape(-1, 3)
+            points = torch.cat((points, extra_positive_points), dim=0)
+            labels = torch.cat((labels, torch.ones((extra_positive_points.shape[0]))))
+
+        if num_negative_extra > 0:
+            neg_idx_list = torch.randperm(len(negative_indices[0]))[:num_negative_extra]
+            extra_negative_points = []
+            for neg_idx in neg_idx_list:
+                extra_negative_points.append([negative_indices[i][neg_idx] for i in range(spacial_dim)])
+            extra_negative_points = torch.tensor(extra_negative_points).reshape(-1, 3)
+            points = torch.cat((points, extra_negative_points), dim=0)
+            labels = torch.cat((labels, torch.zeros((extra_negative_points.shape[0]))))
+            # print('extra_negative_points ', extra_negative_points, extra_negative_points.shape)
+            # print('==> points ', points.shape, labels)
+        
+        if fix_extra_point_num is None:
+            left_point_num = num_positive_extra + num_negative_extra + 1 - labels.shape[0]
+        else:
+            left_point_num = fix_extra_point_num  + 1 - labels.shape[0]
+
+        for _ in range(left_point_num):
+            ignore_point = torch.tensor([-1,-1,-1]).unsqueeze(dim=0)
+            points = torch.cat((points, ignore_point), dim=0)
+            labels = torch.cat((labels, torch.tensor([-1]).reshape(1)))
+
+        return (points, labels)
+
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):
         d = dict(data)
@@ -767,104 +861,6 @@ def _get_scan_interval(
             scan_interval.append(interval if interval > 0 else 1)
     return tuple(scan_interval)
 
-
-def generate_box(pred_pre, bbox_shift=None):
-    meaning_post_label = pred_pre # [h, w, d]
-    ones_idx = (meaning_post_label > 0).nonzero(as_tuple=True)
-    if all(tensor.nelement() == 0 for tensor in ones_idx):
-        bboxes = torch.tensor([-1,-1,-1,-1,-1,-1])
-        # print(bboxes, bboxes.shape)
-        return bboxes
-    min_coords = [dim.min() for dim in ones_idx]    # [x_min, y_min, z_min]
-    max_coords = [dim.max() for dim in ones_idx]    # [x_max, y_max, z_max]
-
-
-    if bbox_shift is None:
-        corner_min = []
-        corner_max = []
-        shape = meaning_post_label.shape
-        for coor in min_coords:
-            coor_ = max(0, coor)
-            corner_min.append(coor_)
-        for idx, coor in enumerate(max_coords):
-            coor_ = min(shape[idx], coor)
-            corner_max.append(coor_)
-        corner_min = torch.tensor(corner_min)
-        corner_max = torch.tensor(corner_max)
-        return torch.cat((corner_min, corner_max), dim=0)
-    else:
-        # add perturbation to bounding box coordinates
-        corner_min = []
-        corner_max = []
-        shape = meaning_post_label.shape
-        for coor in min_coords:
-            coor_ = max(0, coor + random.randint(-bbox_shift, bbox_shift))
-            corner_min.append(coor_)
-        for idx, coor in enumerate(max_coords):
-            coor_ = min(shape[idx], coor + random.randint(-bbox_shift, bbox_shift))
-            corner_max.append(coor_)
-        corner_min = torch.tensor(corner_min)
-        corner_max = torch.tensor(corner_max)
-        return torch.cat((corner_min, corner_max), dim=0)
-
-
-def select_points(preds, num_positive_extra=4, num_negative_extra=0, fix_extra_point_num=None):
-    spacial_dim = 3
-    points = torch.zeros((0, 3))
-    labels = torch.zeros((0))
-    pos_thred = 0.9
-    neg_thred = 0.1
-    
-    # get pos/net indices
-    positive_indices = torch.nonzero(preds > pos_thred, as_tuple=True) # ([pos x], [pos y], [pos z])
-    negative_indices = torch.nonzero(preds < neg_thred, as_tuple=True)
-
-    ones_idx = (preds > pos_thred).nonzero(as_tuple=True)
-    if all(tmp.nelement() == 0 for tmp in ones_idx):
-        # all neg
-        num_positive_extra = 0
-        selected_positive_point = torch.tensor([-1,-1,-1]).unsqueeze(dim=0)
-        points = torch.cat((points, selected_positive_point), dim=0)
-        labels = torch.cat((labels, torch.tensor([-1]).reshape(1)))
-    else:
-        # random select a pos point
-        random_idx = torch.randint(len(positive_indices[0]), (1,))
-        selected_positive_point = torch.tensor([positive_indices[i][random_idx] for i in range(spacial_dim)]).unsqueeze(dim=0)
-        points = torch.cat((points, selected_positive_point), dim=0)
-        labels = torch.cat((labels, torch.ones((1))))
-
-    if num_positive_extra > 0:
-        pos_idx_list = torch.randperm(len(positive_indices[0]))[:num_positive_extra]
-        extra_positive_points = []
-        for pos_idx in pos_idx_list:
-            extra_positive_points.append([positive_indices[i][pos_idx] for i in range(spacial_dim)])
-        extra_positive_points = torch.tensor(extra_positive_points).reshape(-1, 3)
-        points = torch.cat((points, extra_positive_points), dim=0)
-        labels = torch.cat((labels, torch.ones((extra_positive_points.shape[0]))))
-
-    if num_negative_extra > 0:
-        neg_idx_list = torch.randperm(len(negative_indices[0]))[:num_negative_extra]
-        extra_negative_points = []
-        for neg_idx in neg_idx_list:
-            extra_negative_points.append([negative_indices[i][neg_idx] for i in range(spacial_dim)])
-        extra_negative_points = torch.tensor(extra_negative_points).reshape(-1, 3)
-        points = torch.cat((points, extra_negative_points), dim=0)
-        labels = torch.cat((labels, torch.zeros((extra_negative_points.shape[0]))))
-        # print('extra_negative_points ', extra_negative_points, extra_negative_points.shape)
-        # print('==> points ', points.shape, labels)
-    
-    if fix_extra_point_num is None:
-        left_point_num = num_positive_extra + num_negative_extra + 1 - labels.shape[0]
-    else:
-        left_point_num = fix_extra_point_num  + 1 - labels.shape[0]
-
-    for _ in range(left_point_num):
-        ignore_point = torch.tensor([-1,-1,-1]).unsqueeze(dim=0)
-        points = torch.cat((points, ignore_point), dim=0)
-        labels = torch.cat((labels, torch.tensor([-1]).reshape(1)))
-
-    return (points, labels)
-
 # build 3D SAM
 import torch
 import numpy as np