add model

model_segvol_single.py

@@ -123,7 +123,7 @@ class SegVolModel(PreTrainedModel):
 class SegVolProcessor():
     def __init__(self, spatial_size) -> None:
         self.img_loader = transforms.LoadImage()
-        self.transform = transforms.Compose(
+        self.transform4test = transforms.Compose(
             [
                 ForegroundNormalization(keys=["image"]),
                 DimTranspose(keys=["image", "label"]),
@@ -134,6 +134,36 @@ 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')
+        self.transform4train = transforms.Compose(
+        [
+            transforms.AddChanneld(keys=["image"]),
+            DimTranspose(keys=["image", "label"]),
+            MinMaxNormalization(),
+            transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
+            transforms.SpatialPadd(keys=["image", "label"], spatial_size=spatial_size, mode='constant'),
+            transforms.OneOf(transforms=[
+                transforms.Resized(keys=["image", "label"],spatial_size=spatial_size),
+                transforms.RandCropByPosNegLabeld(
+                    keys=["image", "label"],
+                    label_key="label",
+                    spatial_size=spatial_size,
+                    pos=5,
+                    neg=1,
+                    num_samples=1,
+                    image_key="image",
+                    image_threshold=0,
+                ),
+                ],
+                weights=[1, 3]
+            ),
+            transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=0),
+            transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=1),
+            transforms.RandFlipd(keys=["image", "label"], prob=0.2, spatial_axis=2),
+            transforms.RandScaleIntensityd(keys="image", factors=0.2, prob=0.2),
+            transforms.RandShiftIntensityd(keys="image", offsets=0.2, prob=0.2),
+            transforms.ToTensord(keys=["image", "label"]),
+        ]
+    )
     
     # 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
@@ -174,7 +204,7 @@ class SegVolProcessor():
             'image': ct_npy,
             'label': gt_npy
         }
-        item = self.transform(item)
+        item = self.transform4test(item)
         item_zoom_out = self.zoom_out_transform(item)
         item['zoom_out_image'] = item_zoom_out['image']
         item['zoom_out_label'] = item_zoom_out['label']
@@ -223,6 +253,17 @@ class SegVolProcessor():
         preds_nii = nib.Nifti1Image(preds_save, affine=ct.affine, header=ct.header)
         nib.save(preds_nii, save_path)
 
+    def train_transform(self, ct_npy, gt_npy):
+        item = {
+            'image': ct_npy,
+            'label': gt_npy
+        }
+        item = self.transform4train(item)
+        if type(item) is list:
+            assert len(item) == 1
+            item = item[0]
+        return item
+
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):
         d = dict(data)
@@ -409,8 +450,8 @@ class SegVol(nn.Module):
         ## sl
         sl_loss = self.supervised_forward(image, image_embedding, img_shape, kwargs['train_organs'], kwargs['train_labels'])
         ## ssl
-        ssl_loss = self.unsupervised_forward(image, image_embedding, kwargs['pseudo_seg_cleaned'], img_shape)
-        return sl_loss, ssl_loss
+        # ssl_loss = self.unsupervised_forward(image, image_embedding, kwargs['pseudo_seg_cleaned'], img_shape)
+        return sl_loss
 
     def forward_decoder(self, image_embedding, img_shape, text=None, boxes=None, points=None):
         with torch.no_grad():
@@ -456,20 +497,20 @@ class SegVol(nn.Module):
             sl_loss += sl_loss_dice + sl_loss_bce
         return sl_loss
     
-    def unsupervised_forward(self, image, image_embedding, pseudo_seg_cleaned, img_shape):
-        sll_loss = 0
-        for iter in range(self.decoder_iter):
-            if iter % 2 == 0:
-                pseudo_labels, pseudo_points_prompt = self.build_pseudo_point_prompt_label(image.shape, pseudo_seg_cleaned)
-                logits = self.forward_decoder(image_embedding, img_shape, text=None, boxes=None, points=pseudo_points_prompt)
-            else:
-                pseudo_labels, pseudo_bboxes_prompt = self.build_pseudo_box_prompt_label(image.shape, pseudo_seg_cleaned)
-                logits = self.forward_decoder(image_embedding, img_shape, text=None, boxes=pseudo_bboxes_prompt, points=None)
-            # cal loss
-            sll_loss_dice = self.dice_loss.forward(logits.squeeze().float(), pseudo_labels.squeeze().float())
-            sll_loss_bce = self.bce_loss.forward(logits.squeeze().float(), pseudo_labels.squeeze().float())
-            sll_loss += sll_loss_dice + sll_loss_bce
-        return sll_loss
+    # def unsupervised_forward(self, image, image_embedding, pseudo_seg_cleaned, img_shape):
+    #     sll_loss = 0
+    #     for iter in range(self.decoder_iter):
+    #         if iter % 2 == 0:
+    #             pseudo_labels, pseudo_points_prompt = self.build_pseudo_point_prompt_label(image.shape, pseudo_seg_cleaned)
+    #             logits = self.forward_decoder(image_embedding, img_shape, text=None, boxes=None, points=pseudo_points_prompt)
+    #         else:
+    #             pseudo_labels, pseudo_bboxes_prompt = self.build_pseudo_box_prompt_label(image.shape, pseudo_seg_cleaned)
+    #             logits = self.forward_decoder(image_embedding, img_shape, text=None, boxes=pseudo_bboxes_prompt, points=None)
+    #         # cal loss
+    #         sll_loss_dice = self.dice_loss.forward(logits.squeeze().float(), pseudo_labels.squeeze().float())
+    #         sll_loss_bce = self.bce_loss.forward(logits.squeeze().float(), pseudo_labels.squeeze().float())
+    #         sll_loss += sll_loss_dice + sll_loss_bce
+    #     return sll_loss
 
     def build_prompt_label(self, bs, training_organs, train_labels):
         # generate prompt & label
@@ -501,68 +542,68 @@ class SegVol(nn.Module):
         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).to(self.custom_device)
-        # generate points
-        points = []
-        point_labels = []
-        for batch_idx in range(input_shape[0]):
-            # generate pseudo label
-            unique_ids = torch.unique(seg_labels[batch_idx])
-            unique_ids = unique_ids[unique_ids != -1]
-            region_id = random.choice(unique_ids).item()
-            pseudo_labels[batch_idx][seg_labels[batch_idx]==region_id] = 1
-            # generate point prompt
-            num_positive_extra_max, num_negative_extra_max = 10, 10
-            num_positive_extra = random.randint(4, num_positive_extra_max)
-            num_negative_extra = random.randint(0, num_negative_extra_max)
-            assert len(pseudo_labels[batch_idx][0].shape) == 3
-            point, point_label = select_points(
-                pseudo_labels[batch_idx][0],
-                num_positive_extra=num_positive_extra,
-                num_negative_extra=num_negative_extra,
-                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).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).to(self.custom_device)
-        iter_bboxes = []
-        # generate boxes
-        for batch_idx in range(input_shape[0]):
-            # generate ori pseudo label
-            unique_ids = torch.unique(seg_labels_cleaned[batch_idx])
-            unique_ids = unique_ids[unique_ids != -1]
-            region_id = random.choice(unique_ids).item()
-            pseudo_labels[batch_idx][seg_labels_cleaned[batch_idx]==region_id] = 1
-            # generate box prompt
-            box = generate_box(pseudo_labels[batch_idx][0])
-            iter_bboxes.append(box)
-            # refine pseudo label
-            x_min, y_min, z_min, x_max, y_max, z_max = box
-            binary_cube = torch.zeros_like(pseudo_labels[batch_idx][0]).int()
-            binary_cube[x_min:x_max+1, y_min:y_max+1, z_min:z_max+1] = 1
-            # cal iou
-            mask_label = seg_labels_cleaned[batch_idx][0]
-            assert binary_cube.shape == mask_label.shape, str(binary_cube.shape) + ' ' + str(mask_label.shape)
-            mask_values_in_binary_cube = mask_label[binary_cube == 1]
-            unique_mask_values = torch.unique(mask_values_in_binary_cube)
-            # print('unique_mask_values ', unique_mask_values)
-            for value in unique_mask_values:
-                if value == -1: continue
-                mask_area = (mask_label == value)
-                intersection = (binary_cube & mask_area)
-                iou = intersection.float().sum() / mask_area.float().sum()
-                if iou > 0.90:
-                    # 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().to(self.custom_device)
-        return pseudo_labels, bboxes
+    # def build_pseudo_point_prompt_label(self, input_shape, seg_labels):
+    #     pseudo_labels = torch.zeros(input_shape).to(self.custom_device)
+    #     # generate points
+    #     points = []
+    #     point_labels = []
+    #     for batch_idx in range(input_shape[0]):
+    #         # generate pseudo label
+    #         unique_ids = torch.unique(seg_labels[batch_idx])
+    #         unique_ids = unique_ids[unique_ids != -1]
+    #         region_id = random.choice(unique_ids).item()
+    #         pseudo_labels[batch_idx][seg_labels[batch_idx]==region_id] = 1
+    #         # generate point prompt
+    #         num_positive_extra_max, num_negative_extra_max = 10, 10
+    #         num_positive_extra = random.randint(4, num_positive_extra_max)
+    #         num_negative_extra = random.randint(0, num_negative_extra_max)
+    #         assert len(pseudo_labels[batch_idx][0].shape) == 3
+    #         point, point_label = select_points(
+    #             pseudo_labels[batch_idx][0],
+    #             num_positive_extra=num_positive_extra,
+    #             num_negative_extra=num_negative_extra,
+    #             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).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).to(self.custom_device)
+    #     iter_bboxes = []
+    #     # generate boxes
+    #     for batch_idx in range(input_shape[0]):
+    #         # generate ori pseudo label
+    #         unique_ids = torch.unique(seg_labels_cleaned[batch_idx])
+    #         unique_ids = unique_ids[unique_ids != -1]
+    #         region_id = random.choice(unique_ids).item()
+    #         pseudo_labels[batch_idx][seg_labels_cleaned[batch_idx]==region_id] = 1
+    #         # generate box prompt
+    #         box = generate_box(pseudo_labels[batch_idx][0])
+    #         iter_bboxes.append(box)
+    #         # refine pseudo label
+    #         x_min, y_min, z_min, x_max, y_max, z_max = box
+    #         binary_cube = torch.zeros_like(pseudo_labels[batch_idx][0]).int()
+    #         binary_cube[x_min:x_max+1, y_min:y_max+1, z_min:z_max+1] = 1
+    #         # cal iou
+    #         mask_label = seg_labels_cleaned[batch_idx][0]
+    #         assert binary_cube.shape == mask_label.shape, str(binary_cube.shape) + ' ' + str(mask_label.shape)
+    #         mask_values_in_binary_cube = mask_label[binary_cube == 1]
+    #         unique_mask_values = torch.unique(mask_values_in_binary_cube)
+    #         # print('unique_mask_values ', unique_mask_values)
+    #         for value in unique_mask_values:
+    #             if value == -1: continue
+    #             mask_area = (mask_label == value)
+    #             intersection = (binary_cube & mask_area)
+    #             iou = intersection.float().sum() / mask_area.float().sum()
+    #             if iou > 0.90:
+    #                 # 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().to(self.custom_device)
+    #     return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
     def __init__(self, clip_model):