add processor

config_segvol.py

@@ -5,9 +5,12 @@ class SegVolConfig(PretrainedConfig):
 
     def __init__(
         self,
+        test_mode=True,
+        test_w_zoom=False,
         **kwargs,
     ):
         self.spatial_size = [32, 256, 256]
         self.patch_size = [4, 16, 16]
-        self.test_mode = True
+        self.test_mode = test_mode
+        self.test_w_zoom = test_w_zoom
         super().__init__(**kwargs)

model_segvol_single.py

@@ -1,5 +1,7 @@
 from transformers import PreTrainedModel
 from .config_segvol import SegVolConfig
+import numpy as np
+import monai.transforms as transforms
 
 class SegVolModel(PreTrainedModel):
     config_class = SegVolConfig
@@ -21,10 +23,114 @@ class SegVolModel(PreTrainedModel):
             patch_size=self.config.patch_size,
             test_mode=self.config.test_mode,
             )
+        
+        self.processor = SegVolProcessor(spatial_size=self.config.spatial_size)
 
     def forward(self, image, text=None, boxes=None, points=None, **kwargs):
         return self.model.forward(image, text=text, boxes=boxes, points=points, **kwargs)
 
+# processor
+class SegVolProcessor():
+    def __init__(self, spatial_size) -> None:
+        self.img_loader = transforms.LoadImage()
+        self.transform = transforms.Compose(
+            [
+                ForegroundNormalization(keys=["image"]),
+                DimTranspose(keys=["image", "label"]),
+                MinMaxNormalization(),
+                transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
+                transforms.ToTensord(keys=["image", "label"]),
+            ]
+        )
+        self.zoom_out_transform = transforms.Resized(keys=["image", "label"], spatial_size=spatial_size, mode='nearest-exact')
+    
+    # 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
+    def preprocess_ct_gt(self, ct_path, gt_path, category):
+        item = {}
+        # generate ct_voxel_ndarray
+        ct_voxel_ndarray, _ = self.img_loader(ct_path)
+        ct_voxel_ndarray = np.array(ct_voxel_ndarray).squeeze()
+        ct_shape = ct_voxel_ndarray.shape
+        ct_voxel_ndarray = np.expand_dims(ct_voxel_ndarray, axis=0)
+        item['image'] = ct_voxel_ndarray
+
+        # generate gt_voxel_ndarray
+        gt_voxel_ndarray, _ = self.img_loader(gt_path)
+        gt_voxel_ndarray = np.array(gt_voxel_ndarray)
+        present_categories = np.unique(gt_voxel_ndarray)
+        gt_masks = []
+        for cls_idx in range(len(category)):
+            # ignore background
+            cls = cls_idx + 1
+            if cls not in present_categories:
+                gt_voxel_ndarray_category = np.zeros(ct_shape)
+                gt_masks.append(gt_voxel_ndarray_category)
+            else:
+                gt_voxel_ndarray_category = gt_voxel_ndarray.copy()
+                gt_voxel_ndarray_category[gt_voxel_ndarray != cls] = 0
+                gt_voxel_ndarray_category[gt_voxel_ndarray == cls] = 1
+                gt_masks.append(gt_voxel_ndarray_category)
+        gt_voxel_ndarray = np.stack(gt_masks, axis=0)
+        assert gt_voxel_ndarray.shape[0] == len(category) and gt_voxel_ndarray.shape[1:] == ct_voxel_ndarray.shape[1:]
+        item['label'] = gt_voxel_ndarray.astype(np.int32)
+
+        # 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,
+            )
+        return item
+
+class MinMaxNormalization(transforms.Transform):
+    def __call__(self, data):
+        d = dict(data)
+        k = "image"
+        d[k] = d[k] - d[k].min()
+        d[k] = d[k] / np.clip(d[k].max(), a_min=1e-8, a_max=None)
+        return d
+
+class DimTranspose(transforms.Transform):
+    def __init__(self, keys):
+        self.keys = keys
+    
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            d[key] = np.swapaxes(d[key], -1, -3)
+        return d
+
+class ForegroundNormalization(transforms.Transform):
+    def __init__(self, keys):
+        self.keys = keys
+    
+    def __call__(self, data):
+        d = dict(data)
+        
+        for key in self.keys:
+            d[key] = self.normalize(d[key])
+        return d
+    
+    def normalize(self, ct_narray):
+        ct_voxel_ndarray = ct_narray.copy()
+        ct_voxel_ndarray = ct_voxel_ndarray.flatten()
+        thred = np.mean(ct_voxel_ndarray)
+        voxel_filtered = ct_voxel_ndarray[(ct_voxel_ndarray > thred)]
+        upper_bound = np.percentile(voxel_filtered, 99.95)
+        lower_bound = np.percentile(voxel_filtered, 00.05)
+        mean = np.mean(voxel_filtered)
+        std = np.std(voxel_filtered)
+        ### transform ###
+        ct_narray = np.clip(ct_narray, lower_bound, upper_bound)
+        ct_narray = (ct_narray - mean) / max(std, 1e-8)
+        return ct_narray
+    
 # SegVol
 import torch
 import torch.nn as nn