update to support metatensor (#826)

Signed-off-by: Yiheng Wang <[email protected]>

Author: yiheng-wang-nv

kaggle/RANZCR/4th_place_solution/configs/cfg_seg_40_1024d_full.py

@@ -1,9 +1,7 @@
 import numpy as np
 from monai.transforms import (
-    CastToTyped,
     CenterSpatialCropd,
     Compose,
-    EnsureTyped,
     Lambdad,
     NormalizeIntensityd,
     RandAffined,
@@ -60,10 +58,8 @@
             max_spatial_size=(84, 84),
             prob=0.5,
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -80,10 +76,8 @@
         CenterSpatialCropd(
             keys=("input", "mask"), roi_size=(cfg.img_size[0], cfg.img_size[1])
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -97,9 +91,7 @@
             align_corners=False,
         ),
         SpatialPadd(keys=("input", "mask"), spatial_size=(1008, 1008)),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )

kaggle/RANZCR/4th_place_solution/configs/cfg_seg_philipp_16_ch_1024_ap_full.py

@@ -1,9 +1,7 @@
 import numpy as np
 from monai.transforms import (
-    CastToTyped,
     CenterSpatialCropd,
     Compose,
-    EnsureTyped,
     Lambdad,
     NormalizeIntensityd,
     RandAffined,
@@ -62,10 +60,8 @@
             max_spatial_size=(102, 102),
             prob=0.5,
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -82,10 +78,8 @@
         CenterSpatialCropd(
             keys=("input", "mask"), roi_size=(cfg.img_size[0], cfg.img_size[1])
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -99,9 +93,7 @@
             align_corners=False,
         ),
         SpatialPadd(keys=("input", "mask"), spatial_size=(1120, 1120)),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )

kaggle/RANZCR/4th_place_solution/configs/cfg_seg_philipp_16_ch_1024_nons_full.py

@@ -1,9 +1,7 @@
 import numpy as np
 from monai.transforms import (
-    CastToTyped,
     CenterSpatialCropd,
     Compose,
-    EnsureTyped,
     Lambdad,
     NormalizeIntensityd,
     RandAffined,
@@ -63,10 +61,8 @@
             max_spatial_size=(102, 102),
             prob=0.5,
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -83,10 +79,8 @@
         CenterSpatialCropd(
             keys=("input", "mask"), roi_size=(cfg.img_size[0], cfg.img_size[1])
         ),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )
 
@@ -100,9 +94,7 @@
             align_corners=False,
         ),
         SpatialPadd(keys=("input", "mask"), spatial_size=(1120, 1120)),
-        CastToTyped(keys="input", dtype=np.float32),
         NormalizeIntensityd(keys="input", nonzero=False),
         Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
-        EnsureTyped(keys=("input", "mask")),
     ]
 )

kaggle/RANZCR/4th_place_solution/data/seg_data.py

@@ -72,7 +72,7 @@ def load_one(self, study_id: str):
         """
         ext = self.cfg.image_extension
         fp = self.data_folder + study_id + ext
-        img = self.img_reader(filename=fp).transpose(1, 0)
+        img = self.img_reader(filename=fp).numpy().transpose(1, 0)
         img = img[:, :, None]
 
         return img
@@ -129,7 +129,7 @@ def __getitem__(self, idx):
         img = self.load_one(study_id)
         # convert the shape into (Channel, height, width)
         mask = self.get_mask(study_id, img.shape, is_annotated).transpose(2, 0, 1)
-        data = {"input": img.transpose(2, 0, 1), "mask": mask}
+        data = {"input": torch.tensor(img.transpose(2, 0, 1)), "mask": torch.tensor(mask)}
         if self.aug:
             data = self.aug(data)
 

kaggle/RANZCR/4th_place_solution/train.py

@@ -224,8 +224,6 @@ def run_eval(model, val_dataloader, cfg, writer, epoch):
 
         val_preds = torch.cat(val_preds)
         val_targets = torch.cat(val_targets)
-        val_preds = val_preds.cpu().numpy().astype(np.float32)
-        val_targets = val_targets.cpu().numpy().astype(np.float32)
         avg_auc = compute_roc_auc(val_preds, val_targets, average="macro")
         writer.add_scalar("val_avg_auc", avg_auc, epoch)
 

modules/dynunet_pipeline/transforms.py

@@ -4,7 +4,7 @@
                               NormalizeIntensity, RandCropByPosNegLabeld,
                               RandFlipd, RandGaussianNoised,
                               RandGaussianSmoothd, RandScaleIntensityd,
-                              RandZoomd, SpatialCrop, SpatialPadd, EnsureTyped)
+                              RandZoomd, SpatialCrop, SpatialPadd, ToTensord, EnsureTyped)
 from monai.transforms.compose import MapTransform
 from monai.transforms.utils import generate_spatial_bounding_box
 from skimage.transform import resize
@@ -31,6 +31,7 @@ def get_task_transforms(mode, task_id, pos_sample_num, neg_sample_num, num_sampl
             normalize_values=normalize_values[task_id],
             model_mode=mode,
         ),
+        ToTensord(keys="image"),
     ]
     # 3. spatial transforms
     if mode == "train":
@@ -66,23 +67,11 @@ def get_task_transforms(mode, task_id, pos_sample_num, neg_sample_num, num_sampl
             RandFlipd(["image", "label"], spatial_axis=[0], prob=0.5),
             RandFlipd(["image", "label"], spatial_axis=[1], prob=0.5),
             RandFlipd(["image", "label"], spatial_axis=[2], prob=0.5),
-            CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
-            EnsureTyped(keys=["image", "label"]),
-        ]
-    elif mode == "validation":
-        other_transforms = [
-            CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
-            EnsureTyped(keys=["image", "label"]),
-        ]
-    else:
-        other_transforms = [
-            CastToTyped(keys=["image"], dtype=(np.float32)),
-            EnsureTyped(keys=["image"]),
         ]
 
-    all_transforms = load_transforms + sample_transforms + other_transforms
-    return Compose(all_transforms)
-
+        return Compose(load_transforms + sample_transforms + other_transforms)
+    else:
+        return Compose(load_transforms + sample_transforms)
 
 def resample_image(image, shape, anisotrophy_flag):
     resized_channels = []
@@ -273,6 +262,7 @@ def __call__(self, data):
         # load data
         d = dict(data)
         image = d["image"]
+
         image_spacings = d["image_meta_dict"]["pixdim"][1:4].tolist()
 
         if "label" in self.keys:
@@ -294,6 +284,8 @@ def __call__(self, data):
         # calculate shape
         resample_flag = False
         anisotrophy_flag = False
+
+        image = image.numpy()
         if self.target_spacing != image_spacings:
             # resample
             resample_flag = True