ADD: add BraTS segmentation sample application

BraTS sample application addition using BraTS 2020 multi-volume data

Signed-off-by: Cavan Riley <cavan-riley@uiowa.edu>

Author: CavRiley

sample-apps/radiology/lib/configs/segmentation_brats.py

@@ -0,0 +1,107 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+from typing import Any, Dict, Optional, Union
+
+import lib.infers
+import lib.trainers
+from monai.networks.nets import SegResNet
+from monai.utils import optional_import
+
+from monailabel.interfaces.config import TaskConfig
+from monailabel.interfaces.tasks.infer_v2 import InferTask
+from monailabel.interfaces.tasks.train import TrainTask
+from monailabel.utils.others.generic import download_file, remove_file, strtobool
+
+_, has_cp = optional_import("cupy")
+_, has_cucim = optional_import("cucim")
+
+logger = logging.getLogger(__name__)
+
+
+class Segmentation(TaskConfig):
+    def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **kwargs):
+        super().init(name, model_dir, conf, planner, **kwargs)
+
+        # BraTS labels: 3 multi-label channels produced by ConvertToMultiChannelBasedOnBratsClassesd
+        #   Channel 0: TC  - Tumor Core        (label 2 OR label 3)
+        #   Channel 1: WT  - Whole Tumor       (label 1 OR label 2 OR label 3)
+        #   Channel 2: ET  - Enhancing Tumor   (label 2)
+        self.labels = {
+            "tumor core": 1,  # Tumor Core
+            "whole tumor": 2,  # Whole Tumor
+            "enhancing tumor": 3,  # Enhancing Tumor
+        }
+
+        # Model Files
+        self.path = [
+            os.path.join(self.model_dir, f"pretrained_{name}.pt"),  # pretrained
+            os.path.join(self.model_dir, f"{name}.pt"),  # published
+        ]
+
+        # Download PreTrained Model (optional)
+        if strtobool(self.conf.get("use_pretrained_model", "false")):
+            url = f"{self.conf.get('pretrained_path', self.PRE_TRAINED_PATH)}"
+            url = f"{url}/radiology_segmentation_segresnet_brats.pt"
+            download_file(url, self.path[0])
+
+        # Spacing and ROI for BraTS (isotropic 1mm, large crop matching tutorial)
+        self.target_spacing = (1.0, 1.0, 1.0)
+        self.roi_size = (224, 224, 144)
+
+        # Number of input channels: 4 MRI modalities (FLAIR, T1, T1Gd, T2)
+        # when multi_file=True the LoadDirectoryImagesd loader stacks them;
+        # when multi_file=False the image file must already be a 4-channel volume.
+        try:
+            input_channels = int(self.conf.get("input_channels", 4))
+        except (ValueError, TypeError):
+            logger.warning("Could not parse input_channels, defaulting to 4")
+            input_channels = 4
+
+        # Network
+        self.network = SegResNet(
+            blocks_down=(1, 2, 2, 4),
+            blocks_up=(1, 1, 1),
+            init_filters=16,
+            in_channels=input_channels,
+            out_channels=len(self.labels),  # TC, WT, ET — sigmoid multilabel, no background channel
+            dropout_prob=0.2,
+        )
+
+    def infer(self) -> Union[InferTask, Dict[str, InferTask]]:
+        task: InferTask = lib.infers.Segmentation(
+            path=self.path,
+            network=self.network,
+            roi_size=self.roi_size,
+            target_spacing=self.target_spacing,
+            labels=self.labels,
+            preload=strtobool(self.conf.get("preload", "false")),
+            config={"largest_cc": True if has_cp and has_cucim else False},
+        )
+        return task
+
+    def trainer(self) -> Optional[TrainTask]:
+        output_dir = os.path.join(self.model_dir, self.name)
+        load_path = self.path[0] if os.path.exists(self.path[0]) else self.path[1]
+
+        task: TrainTask = lib.trainers.Segmentation(
+            model_dir=output_dir,
+            network=self.network,
+            roi_size=self.roi_size,
+            target_spacing=self.target_spacing,
+            load_path=load_path,
+            publish_path=self.path[1],
+            description="Train BraTS Segmentation Model (TC/WT/ET multilabel)",
+            labels=self.labels,
+        )
+        return task

sample-apps/radiology/lib/infers/segmentation_brats.py

@@ -0,0 +1,159 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Sequence
+
+from lib.transforms.transforms import ConvertFromMultiChannelBasedOnBratsClassesd, GetCentroidsd, LoadDirectoryImagesd
+from monai.inferers import Inferer, SlidingWindowInferer
+from monai.transforms import (
+    Activationsd,
+    AsDiscreted,
+    EnsureChannelFirstd,
+    EnsureTyped,
+    KeepLargestConnectedComponentd,
+    LoadImaged,
+    NormalizeIntensityd,
+    Orientationd,
+    Spacingd,
+)
+
+from monailabel.interfaces.tasks.infer_v2 import InferType
+from monailabel.tasks.infer.basic_infer import BasicInferTask
+from monailabel.transform.post import Restored
+
+
+class Segmentation(BasicInferTask):
+    """
+    Inference Engine for BraTS brain tumour segmentation using a SegResNet.
+
+    The model outputs 3 channels (TC, WT, ET) with sigmoid activations — it is
+    a multilabel task, NOT a softmax classification.  Each channel is thresholded
+    independently at 0.5 to produce binary maps.
+
+    Two image loading modes are supported (set via ``data["multi_file"]``):
+      - False (default): the input image is a single 4-channel NIfTI volume.
+      - True:            ``data["image"]`` is a directory containing 4 single-
+                         modality NIfTI files; LoadDirectoryImagesd stacks them.
+    """
+
+    def __init__(
+        self,
+        path,
+        network=None,
+        target_spacing=(1.0, 1.0, 1.0),
+        type=InferType.SEGMENTATION,
+        labels=None,
+        dimension=3,
+        description="Pre-trained BraTS SegResNet — TC/WT/ET multilabel segmentation",
+        **kwargs,
+    ):
+        super().__init__(
+            path=path,
+            network=network,
+            type=type,
+            labels=labels,
+            dimension=dimension,
+            description=description,
+            load_strict=False,
+            **kwargs,
+        )
+        self.target_spacing = target_spacing
+
+    def pre_transforms(self, data=None) -> Sequence[Callable]:
+        """
+        Pre-processing pipeline matching the official MONAI BraTS tutorial.
+
+        NOTE: ScaleIntensityRangePercentilesd and CenterSpatialCropd from the
+        original file have been removed — they are not part of the BraTS pipeline
+        and would distort MRI intensity normalisation.  NormalizeIntensityd with
+        nonzero=True, channel_wise=True is the correct approach for multi-modal MRI.
+        """
+        channels = data.get("input_channels", 4)
+        t = [
+            (
+                LoadImaged(keys="image", reader="ITKReader", ensure_channel_first=True)
+                if data.get("multi_file", False) is False
+                else LoadDirectoryImagesd(
+                    keys="image",
+                    target_spacing=self.target_spacing,
+                    channels=channels,
+                )
+            ),
+            EnsureTyped(keys="image", device=data.get("device") if data else None),
+            # EnsureChannelFirstd is safe to keep as a guard; if the channel dim is
+            # already present (ITKReader + ensure_channel_first) it is a no-op.
+            EnsureChannelFirstd(keys="image", channel_dim=0),
+            Orientationd(keys="image", axcodes="RAS"),
+            Spacingd(
+                keys="image",
+                pixdim=self.target_spacing,
+                allow_missing_keys=True,
+            ),
+            # Channel-wise intensity normalisation on non-zero voxels only.
+            # This matches both the tutorial and the training pipeline exactly.
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
+        ]
+        return t
+
+    def inferer(self, data=None) -> Inferer:
+        return SlidingWindowInferer(
+            roi_size=self.roi_size,
+            sw_batch_size=2,
+            overlap=0.4,
+            padding_mode="replicate",
+            mode="gaussian",
+        )
+
+    def inverse_transforms(self, data=None):
+        return []
+
+    def post_transforms(self, data=None) -> Sequence[Callable]:
+        """
+        Post-processing for multilabel sigmoid output.
+
+        IMPORTANT differences from a softmax segmentation:
+          - Activationsd uses sigmoid=True (not softmax=True).
+          - AsDiscreted thresholds each channel at 0.5 independently
+            (not argmax, because channels are not mutually exclusive).
+          - KeepLargestConnectedComponentd is applied per-channel if available.
+        """
+        t = [
+            EnsureTyped(keys="pred", device=data.get("device") if data else None),
+            # Sigmoid: each of the 3 channels (TC, WT, ET) is activated independently.
+            Activationsd(keys="pred", sigmoid=True),
+            # Threshold each channel at 0.5 to produce binary masks.
+            AsDiscreted(keys="pred", threshold=0.5),
+        ]
+
+        if data and data.get("largest_cc", False):
+            # Apply per-channel so TC, WT and ET are each cleaned independently.
+            t.append(
+                KeepLargestConnectedComponentd(
+                    keys="pred",
+                    independent=True,  # treat each channel separately
+                )
+            )
+
+        t.extend(
+            [
+                # Merge 3 binary channels → single-channel integer label map
+                # Must happen before Restored so spatial metadata is applied
+                # to the final (1, H, W, D) output, not the intermediate (3, H, W, D).
+                ConvertFromMultiChannelBasedOnBratsClassesd(keys="pred"),
+                Restored(
+                    keys="pred",
+                    ref_image="image",
+                    config_labels=self.labels if data.get("restore_label_idx", False) else None,
+                ),
+                GetCentroidsd(keys="pred", centroids_key="centroids"),
+            ]
+        )
+        return t