Update segmentation_3d tutorials for metatensor (#776)

* Update brats_segmentation_3d tutorial for metatensor

* update spleen_seg_3d for MetaTensor

* update spleen_seg_3d_lightning for MetaTensor

* fix pep8 error for spleen_seg_3d_lightning

* fix transform in brats_seg
update spleen_seg_visualization
update swin_unetr_brats_seg

* update unet_3d catalyst and ignite

* update btcv_seg for MetaTensor

* update challenge baseline and torch folder

* update ignite dir

Co-authored-by: Wenqi Li <831580+wyli@users.noreply.github.com>

Author: KumoLiu

3d_segmentation/brats_segmentation_3d.ipynb

@@ -25,7 +25,7 @@ This directory contains a simple baseline method [using MONAI](https://monai.io)
 
 The script is tested with:
 
-- `Ubuntu 18.04` | `Python 3.6` | `CUDA 10.2`
+- `Ubuntu 20.04` | `Python 3.8` | `CUDA 11.7`
 
 On a GPU with [automatic mixed precision support](https://developer.nvidia.com/automatic-mixed-precision):
 

3d_segmentation/challenge_baseline/README.md

@@ -36,7 +36,6 @@
     ScaleIntensityRanged,
     Spacingd,
     SpatialPadd,
-    EnsureTyped,
 )
 
 
@@ -69,12 +68,12 @@ def get_xforms(mode="train", keys=("image", "label")):
                 RandFlipd(keys, spatial_axis=2, prob=0.5),
             ]
         )
-        dtype = (np.float32, np.uint8)
+        dtype = (torch.float32, torch.uint8)
     if mode == "val":
-        dtype = (np.float32, np.uint8)
+        dtype = (torch.float32, torch.uint8)
     if mode == "infer":
-        dtype = (np.float32,)
-    xforms.extend([CastToTyped(keys, dtype=dtype), EnsureTyped(keys)])
+        dtype = (torch.float32,)
+    xforms.extend([CastToTyped(keys, dtype=dtype)])
     return monai.transforms.Compose(xforms)
 
 
@@ -172,7 +171,7 @@ def train(data_folder=".", model_folder="runs"):
 
     # create evaluator (to be used to measure model quality during training
     val_post_transform = monai.transforms.Compose(
-        [EnsureTyped(keys=("pred", "label")), AsDiscreted(keys=("pred", "label"), argmax=(True, False), to_onehot=2)]
+        [AsDiscreted(keys=("pred", "label"), argmax=(True, False), to_onehot=2)]
     )
     val_handlers = [
         ProgressBar(),
@@ -246,7 +245,7 @@ def infer(data_folder=".", model_folder="runs", prediction_folder="output"):
     saver = monai.data.NiftiSaver(output_dir=prediction_folder, mode="nearest")
     with torch.no_grad():
         for infer_data in infer_loader:
-            logging.info(f"segmenting {infer_data['image_meta_dict']['filename_or_obj']}")
+            logging.info(f"segmenting {infer_data['image'].meta['filename_or_obj']}")
             preds = inferer(infer_data[keys[0]].to(device), net)
             n = 1.0
             for _ in range(4):
@@ -262,7 +261,7 @@ def infer(data_folder=".", model_folder="runs", prediction_folder="output"):
                     n = n + 1.0
             preds = preds / n
             preds = (preds.argmax(dim=1, keepdims=True)).float()
-            saver.save_batch(preds, infer_data["image_meta_dict"])
+            saver.save_batch(preds, infer_data["image"].meta)
 
     # copy the saved segmentations into the required folder structure for submission
     submission_dir = os.path.join(prediction_folder, "to_submit")

3d_segmentation/challenge_baseline/run_net.py

@@ -19,14 +19,13 @@
 import numpy as np
 import torch
 from ignite.engine import Engine
-from torch.utils.data import DataLoader
 
 from monai import config
-from monai.data import ImageDataset, create_test_image_3d, decollate_batch
+from monai.data import ImageDataset, create_test_image_3d, decollate_batch, DataLoader
 from monai.handlers import CheckpointLoader, MeanDice, StatsHandler
 from monai.inferers import sliding_window_inference
 from monai.networks.nets import UNet
-from monai.transforms import Activations, AddChannel, AsDiscrete, Compose, SaveImage, ScaleIntensity, EnsureType
+from monai.transforms import Activations, AddChannel, AsDiscrete, Compose, SaveImage, ScaleIntensity
 
 
 def main(tempdir):
@@ -47,8 +46,8 @@ def main(tempdir):
     segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
 
     # define transforms for image and segmentation
-    imtrans = Compose([ScaleIntensity(), AddChannel(), EnsureType()])
-    segtrans = Compose([AddChannel(), EnsureType()])
+    imtrans = Compose([ScaleIntensity(), AddChannel()])
+    segtrans = Compose([AddChannel()])
     ds = ImageDataset(images, segs, transform=imtrans, seg_transform=segtrans, image_only=False)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -65,7 +64,7 @@ def main(tempdir):
     roi_size = (96, 96, 96)
     sw_batch_size = 4
 
-    post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
     save_image = SaveImage(output_dir="tempdir", output_ext=".nii.gz", output_postfix="seg")
 
     def _sliding_window_processor(engine, batch):
@@ -74,9 +73,8 @@ def _sliding_window_processor(engine, batch):
             val_images, val_labels = batch[0].to(device), batch[1].to(device)
             seg_probs = sliding_window_inference(val_images, roi_size, sw_batch_size, net)
             seg_probs = [post_trans(i) for i in decollate_batch(seg_probs)]
-            val_data = decollate_batch(batch[2])
-            for seg_prob, data in zip(seg_probs, val_data):
-                save_image(seg_prob, data)
+            for seg_prob in seg_probs:
+                save_image(seg_prob)
             return seg_probs, val_labels
 
     evaluator = Engine(_sliding_window_processor)

3d_segmentation/ignite/unet_evaluation_array.py

@@ -26,7 +26,7 @@
 from monai.handlers import CheckpointLoader, MeanDice, StatsHandler
 from monai.inferers import sliding_window_inference
 from monai.networks.nets import UNet
-from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, SaveImage, ScaleIntensityd, EnsureTyped, EnsureType
+from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, SaveImage, ScaleIntensityd
 
 
 def main(tempdir):
@@ -53,7 +53,6 @@ def main(tempdir):
             LoadImaged(keys=["img", "seg"]),
             AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
             ScaleIntensityd(keys="img"),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
     val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
@@ -72,7 +71,7 @@ def main(tempdir):
     roi_size = (96, 96, 96)
     sw_batch_size = 4
 
-    post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
     save_image = SaveImage(output_dir="tempdir", output_ext=".nii.gz", output_postfix="seg")
 
     def _sliding_window_processor(engine, batch):
@@ -81,9 +80,8 @@ def _sliding_window_processor(engine, batch):
             val_images, val_labels = batch["img"].to(device), batch["seg"].to(device)
             seg_probs = sliding_window_inference(val_images, roi_size, sw_batch_size, net)
             seg_probs = [post_trans(i) for i in decollate_batch(seg_probs)]
-            val_data = decollate_batch(batch["img_meta_dict"])
-            for seg_prob, data in zip(seg_probs, val_data):
-                save_image(seg_prob, data)
+            for seg_prob in seg_probs:
+                save_image(seg_prob)
             return seg_probs, val_labels
 
     evaluator = Engine(_sliding_window_processor)

3d_segmentation/ignite/unet_evaluation_dict.py

@@ -20,10 +20,9 @@
 import torch
 from ignite.engine import Events, create_supervised_evaluator, create_supervised_trainer
 from ignite.handlers import EarlyStopping, ModelCheckpoint
-from torch.utils.data import DataLoader
 
 import monai
-from monai.data import ImageDataset, create_test_image_3d, decollate_batch
+from monai.data import ImageDataset, create_test_image_3d, decollate_batch, DataLoader
 from monai.handlers import (
     MeanDice,
     StatsHandler,
@@ -39,7 +38,6 @@
     RandSpatialCrop,
     Resize,
     ScaleIntensity,
-    EnsureType,
 )
 
 
@@ -67,16 +65,15 @@ def main(tempdir):
             ScaleIntensity(),
             AddChannel(),
             RandSpatialCrop((96, 96, 96), random_size=False),
-            EnsureType(),
         ]
     )
     train_segtrans = Compose(
-        [AddChannel(), RandSpatialCrop((96, 96, 96), random_size=False), EnsureType()]
+        [AddChannel(), RandSpatialCrop((96, 96, 96), random_size=False)]
     )
     val_imtrans = Compose(
-        [ScaleIntensity(), AddChannel(), Resize((96, 96, 96)), EnsureType()]
+        [ScaleIntensity(), AddChannel(), Resize((96, 96, 96))]
     )
-    val_segtrans = Compose([AddChannel(), Resize((96, 96, 96)), EnsureType()])
+    val_segtrans = Compose([AddChannel(), Resize((96, 96, 96))])
 
     # define image dataset, data loader
     check_ds = ImageDataset(
@@ -151,8 +148,8 @@ def main(tempdir):
     # add evaluation metric to the evaluator engine
     val_metrics = {metric_name: MeanDice()}
 
-    post_pred = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
-    post_label = Compose([EnsureType(), AsDiscrete(threshold=0.5)])
+    post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_label = Compose([AsDiscrete(threshold=0.5)])
 
     # Ignite evaluator expects batch=(img, seg) and returns output=(y_pred, y) at every iteration,
     # user can add output_transform to return other values

3d_segmentation/ignite/unet_training_array.py

@@ -45,8 +45,6 @@
     RandCropByPosNegLabeld,
     RandRotate90d,
     ScaleIntensityd,
-    EnsureTyped,
-    EnsureType,
 )
 
 
@@ -85,15 +83,13 @@ def main(tempdir):
                 num_samples=4,
             ),
             RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 2]),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
     val_transforms = Compose(
         [
             LoadImaged(keys=["img", "seg"]),
             AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
             ScaleIntensityd(keys="img"),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 
@@ -180,8 +176,8 @@ def prepare_batch(batch, device=None, non_blocking=False):
     # add evaluation metric to the evaluator engine
     val_metrics = {metric_name: MeanDice()}
 
-    post_pred = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
-    post_label = Compose([EnsureType(), AsDiscrete(threshold=0.5)])
+    post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_label = Compose([AsDiscrete(threshold=0.5)])
 
     # Ignite evaluator expects batch=(img, seg) and returns output=(y_pred, y) at every iteration,
     # user can add output_transform to return other values