add model

model_segvol_single.py

@@ -26,6 +26,8 @@ class SegVolModel(PreTrainedModel):
         
         self.processor = SegVolProcessor(spatial_size=self.config.spatial_size)
 
+        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):
         print(image.shape, zoomed_image.shape, text_prompt)
         print(bbox_prompt[0].shape, bbox_prompt[1].shape, point_prompt[0].shape, point_prompt[1].shape)
@@ -51,6 +53,7 @@ class SegVolProcessor():
             ]
         )
         self.zoom_out_transform = transforms.Resized(keys=["image", "label"], spatial_size=spatial_size, mode='nearest-exact')
+        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
     # ct_path is path for a ct scan file with nii.gz format
     # gt_path is path for a ground truth file with nii.gz format
@@ -97,15 +100,14 @@ class SegVolProcessor():
         item['zoom_out_label'] = item_zoom_out['label']
         return item
 
-    def point_prompt_b(self, label_single_resize):
-        print(label_single_resize.shape, torch.unique(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()) 
+    def point_prompt_b(self, label_single_resize, num_positive_extra=4, num_negative_extra=0):
+        point, point_label = select_points(label_single_resize, num_positive_extra=num_positive_extra, num_negative_extra=num_negative_extra)
+        points_single = (point.unsqueeze(0).float().to(self.custom_device), point_label.unsqueeze(0).float().to(self.custom_device)) 
         binary_points_resize = build_binary_points(point, point_label, label_single_resize.shape).unsqueeze(0).unsqueeze(0)
         return points_single, binary_points_resize
     
     def bbox_prompt_b(self, label_single_resize):
-        box_single = generate_box(label_single_resize).unsqueeze(0).float().cuda()
+        box_single = generate_box(label_single_resize).unsqueeze(0).float().to(self.custom_device)
         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
 
@@ -269,14 +271,15 @@ class SegVol(nn.Module):
                 test_mode=False,
                 ):
         super().__init__()
+        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.image_encoder = image_encoder
         self.mask_decoder = mask_decoder
         self.prompt_encoder = prompt_encoder
         self.text_encoder = TextEncoder()
         self.feat_shape = np.array(roi_size)/np.array(patch_size)
         self.test_mode = test_mode
-        self.dice_loss = BinaryDiceLoss().cuda()
-        self.bce_loss = BCELoss().cuda()
+        self.dice_loss = BinaryDiceLoss().to(self.custom_device)
+        self.bce_loss = BCELoss().to(self.custom_device)
         self.decoder_iter = 6
 
     def forward(self, image, text=None, boxes=None, points=None, **kwargs):
@@ -379,14 +382,14 @@ class SegVol(nn.Module):
             iter_points_ax.append(point)
             iter_point_labels.append(point_label)
         # batched prompt
-        iter_points_ax = torch.stack(iter_points_ax, dim=0).cuda()
-        iter_point_labels = torch.stack(iter_point_labels, dim=0).cuda()
+        iter_points_ax = torch.stack(iter_points_ax, dim=0).to(self.custom_device)
+        iter_point_labels = torch.stack(iter_point_labels, dim=0).to(self.custom_device)
         iter_points = (iter_points_ax, iter_point_labels)
-        iter_bboxes = torch.stack(iter_bboxes, dim=0).float().cuda()
+        iter_bboxes = torch.stack(iter_bboxes, dim=0).float().to(self.custom_device)
         return iter_points, iter_bboxes, iter_organs
     
     def build_pseudo_point_prompt_label(self, input_shape, seg_labels):
-        pseudo_labels = torch.zeros(input_shape).cuda()
+        pseudo_labels = torch.zeros(input_shape).to(self.custom_device)
         # generate points
         points = []
         point_labels = []
@@ -408,13 +411,13 @@ class SegVol(nn.Module):
                 fix_extra_point_num=num_positive_extra_max + num_negative_extra_max)
             points.append(point)
             point_labels.append(point_label)
-        points = torch.stack(points, dim=0).cuda()
-        point_labels = torch.stack(point_labels, dim=0).cuda()
+        points = torch.stack(points, dim=0).to(self.custom_device)
+        point_labels = torch.stack(point_labels, dim=0).to(self.custom_device)
         pseudo_points_prompt = (points, point_labels)
         return pseudo_labels, pseudo_points_prompt
 
     def build_pseudo_box_prompt_label(self, input_shape, seg_labels_cleaned):
-        pseudo_labels = torch.zeros(input_shape).cuda()
+        pseudo_labels = torch.zeros(input_shape).to(self.custom_device)
         iter_bboxes = []
         # generate boxes
         for batch_idx in range(input_shape[0]):
@@ -445,12 +448,13 @@ class SegVol(nn.Module):
                     # print(f"Mask value {value} has IOU > 0.90 in binary cube.")
                     pseudo_labels[batch_idx][seg_labels_cleaned[batch_idx]==value] = 1
 
-        bboxes = torch.stack(iter_bboxes, dim=0).float().cuda()
+        bboxes = torch.stack(iter_bboxes, dim=0).float().to(self.custom_device)
         return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
     def __init__(self):
         super().__init__()
+        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         config = CLIPTextConfig()
         self.clip_text_model = CLIPTextModel(config)
         self.tokenizer = AutoTokenizer.from_pretrained('openai/clip-vit-base-patch32')
@@ -463,7 +467,7 @@ class TextEncoder(nn.Module):
         text_list = ['A computerized tomography of a {}.'.format(organ_name) for organ_name in organ_names]
         tokens = self.tokenizer(text_list, padding=True, return_tensors="pt")
         for key in tokens.keys():
-            tokens[key] = tokens[key].cuda()
+            tokens[key] = tokens[key].to(self.custom_device)
         return tokens
     
     def forward(self, text):
@@ -752,16 +756,16 @@ def sliding_window_inference(
         if use_point:
             window_binary_prompt_map = torch.cat([binary_prompt_map[win_slice] for win_slice in unravel_slice]).to(sw_device)
             point, point_label = select_points(window_binary_prompt_map.squeeze())
-            points = (point.unsqueeze(0).float().cuda(), point_label.unsqueeze(0).float().cuda())  
+            points = (point.unsqueeze(0).float().to(device), point_label.unsqueeze(0).float().to(device))  
             pseudo_label = torch.cat([global_preds[win_slice] for win_slice in unravel_slice]).to(sw_device)
-            boxes = generate_box(pseudo_label.squeeze()).unsqueeze(0).float().cuda()
+            boxes = generate_box(pseudo_label.squeeze()).unsqueeze(0).float().to(device)
         if use_box:
             if num_win == 1:
                 window_binary_prompt_map = torch.cat([binary_prompt_map[win_slice] for win_slice in unravel_slice]).to(sw_device)
-                boxes = generate_box(window_binary_prompt_map.squeeze()).unsqueeze(0).float().cuda()
+                boxes = generate_box(window_binary_prompt_map.squeeze()).unsqueeze(0).float().to(device)
             else:
                 pseudo_label = torch.cat([global_preds[win_slice] for win_slice in unravel_slice]).to(sw_device)
-                boxes = generate_box(pseudo_label.squeeze()).unsqueeze(0).float().cuda()
+                boxes = generate_box(pseudo_label.squeeze()).unsqueeze(0).float().to(device)
         seg_prob_out = predictor(window_data, text, boxes, points)  # batched patch segmentation
         #############
         # convert seg_prob_out to tuple seg_prob_tuple, this does not allocate new memory.