add model

model_segvol_single.py

@@ -89,102 +89,16 @@ class SegVolProcessor():
         item['zoom_out_label'] = item_zoom_out['label']
         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)
+    def point_prompt(self, label_single_resize):
+        point, point_label = select_points(label_single_resize, num_positive_extra=3, num_negative_extra=3)
+        points_single = (point.unsqueeze(0).float().cuda(), point_label.unsqueeze(0).float().cuda()) 
+        binary_points_resize = build_binary_points(point, point_label, label_single_resize.shape)
+        return points_single, binary_points_resize
+    
+    def bbox_prompt(self, label_single_resize):
+        box_single = generate_box(label_single_resize).unsqueeze(0).float().cuda()
+        binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=label_single_resize.shape)
+        return box_single, binary_cube_resize
 
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):
@@ -228,7 +142,105 @@ class ForegroundNormalization(transforms.Transform):
         ct_narray = np.clip(ct_narray, lower_bound, upper_bound)
         ct_narray = (ct_narray - mean) / max(std, 1e-8)
         return ct_narray
+
+# prompts
+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)
+
 # SegVol
 import torch
 import torch.nn as nn