diff --git a/.gitignore b/.gitignore index 0fbd8a8..3251437 100644 --- a/.gitignore +++ b/.gitignore @@ -164,6 +164,7 @@ patch_config.json # Patch-workflow run byproducts (regenerated by the notebooks; e.g. research/*) preprocessed/ cv_results_*/ +inference_patch_config.json # Claude - share skills/settings, keep local config private .claude/settings.local.json diff --git a/fastMONAI/_modidx.py b/fastMONAI/_modidx.py index 9b5b9c8..35fc24b 100644 --- a/fastMONAI/_modidx.py +++ b/fastMONAI/_modidx.py @@ -633,6 +633,8 @@ 'fastMONAI.vision_plot': { 'fastMONAI.vision_plot._get_slice': ('vision_plot.html#_get_slice', 'fastMONAI/vision_plot.py'), 'fastMONAI.vision_plot.find_max_slice': ( 'vision_plot.html#find_max_slice', 'fastMONAI/vision_plot.py'), + 'fastMONAI.vision_plot.show_mask_overlay': ( 'vision_plot.html#show_mask_overlay', + 'fastMONAI/vision_plot.py'), 'fastMONAI.vision_plot.show_med_img': ('vision_plot.html#show_med_img', 'fastMONAI/vision_plot.py'), 'fastMONAI.vision_plot.show_segmentation_comparison': ( 'vision_plot.html#show_segmentation_comparison', 'fastMONAI/vision_plot.py'), diff --git a/fastMONAI/vision_all.py b/fastMONAI/vision_all.py index 62ab1fe..9c78fcc 100644 --- a/fastMONAI/vision_all.py +++ b/fastMONAI/vision_all.py @@ -1,4 +1,5 @@ #imports all of the functions and classes from the fastMONAI library +from .vision_plot import show_segmentation_comparison from .vision_core import * from .vision_data import * from .vision_augmentation import * diff --git a/fastMONAI/vision_plot.py b/fastMONAI/vision_plot.py index 4f887c8..75e50cc 100644 --- a/fastMONAI/vision_plot.py +++ b/fastMONAI/vision_plot.py @@ -1,7 +1,7 @@ # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_vision_plot.ipynb. # %% auto #0 -__all__ = ['validate_anatomical_plane', 'show_med_img', 'find_max_slice', 'show_segmentation_comparison'] +__all__ = ['validate_anatomical_plane', 'show_med_img', 'find_max_slice', 'show_segmentation_comparison', 'show_mask_overlay'] # %% ../nbs/00_vision_plot.ipynb #7ea97ee7-8a55-485d-a75c-ae1a8055f489 import warnings @@ -196,3 +196,57 @@ def show_segmentation_comparison( axes[2].axis('off') plt.tight_layout() + +# %% ../nbs/00_vision_plot.ipynb #72d68ebd +def show_mask_overlay(image, mask, ctx=None, channel: int = 0, slice_index: int = None, + anatomical_plane: int = 2, voxel_size=None, alpha: float = 0.5, + cmap_img: str = 'gray', cmap_mask: str = 'autumn', + ax=None, figsize=None, title=None, **kwargs): + """Overlay `mask` on `image` on a single axis, using the same slicing and orientation as `.show()`. + + Draws the image, then the mask on top with zero voxels transparent, so it lines up with + `MedImage.show()` / `MedMask.show()`. Reuses the same `_get_slice` those methods use. + + Args: + image: Background image (MedImage or object with `.data` [C, H, W, D]). + mask: Overlay mask (MedMask or object with `.data` [C, H, W, D]); zero voxels are transparent. + ctx: Axis to draw on (fastai-style context); same role as `ax`. + channel: Channel to display. + slice_index: Slice to show; if None, uses `find_max_slice(mask, anatomical_plane)`. + anatomical_plane: 0=sagittal, 1=coronal, 2=axial (default). + voxel_size: Voxel spacing for aspect ratio; if None, aspect defaults to 1. + alpha: Overlay opacity. + cmap_img: Colormap for the background image. + cmap_mask: Colormap for the mask overlay. + ax: Axis to draw on (overrides `ctx`). + figsize: Figure size when a new axis is created. + title: Optional axis title. + kwargs: Extra args forwarded to the background image `imshow`. + + Returns: + The matplotlib axis with the overlay. + """ + validate_anatomical_plane(anatomical_plane) + + img_data = image.data if hasattr(image, 'data') else image + mask_data = mask.data if hasattr(mask, 'data') else mask + if hasattr(img_data, 'cpu'): img_data = img_data.cpu() + if hasattr(mask_data, 'cpu'): mask_data = mask_data.cpu() + + if slice_index is None: + slice_index = find_max_slice(mask_data[channel].numpy(), anatomical_plane) + + img_slice, aspect = _get_slice(img_data, channel, slice_index, anatomical_plane, voxel_size) + mask_slice, _ = _get_slice(mask_data, channel, slice_index, anatomical_plane, voxel_size) + + ax = ax if ax is not None else ctx + if ax is None: + _, ax = plt.subplots(figsize=figsize) + + ax.imshow(img_slice, cmap=cmap_img, aspect=aspect, **kwargs) + ax.imshow(np.ma.masked_where(mask_slice == 0, mask_slice), cmap=cmap_mask, alpha=alpha, aspect=aspect) + + if title is not None: + ax.set_title(title) + ax.axis('off') + return ax diff --git a/nbs/00_vision_plot.ipynb b/nbs/00_vision_plot.ipynb index 07b7968..4b23f40 100644 --- a/nbs/00_vision_plot.ipynb +++ b/nbs/00_vision_plot.ipynb @@ -174,6 +174,14 @@ "metadata": {}, "outputs": [], "source": "#| export\ndef show_segmentation_comparison(\n image, ground_truth, prediction,\n slice_index: int = None,\n anatomical_plane: int = 2,\n metric_value: float = None,\n metric_name: str = 'DSC',\n channel: int = 0,\n figsize: tuple = (15, 5),\n cmap_img: str = 'gray',\n cmap_mask: str = 'gray',\n voxel_size = None\n):\n \"\"\"Display 3-panel comparison: Image | Ground Truth | Prediction.\n \n Useful for validating segmentation results, especially after patch-based \n inference where results are not in standard fastai batch format.\n \n Args:\n image: Input image (MedImage, MedMask, or tensor [C, H, W, D])\n ground_truth: Ground truth mask (MedMask or tensor [C, H, W, D])\n prediction: Predicted mask (tensor [C, H, W, D])\n slice_index: Slice to display. If None, uses find_max_slice on ground_truth.\n anatomical_plane: 0=sagittal, 1=coronal, 2=axial (default)\n metric_value: Optional metric value to display in prediction title\n metric_name: Name of metric for title (default 'DSC')\n channel: Channel to display for multi-channel data (default 0)\n figsize: Figure size (default (15, 5))\n cmap_img: Colormap for image (default 'gray')\n cmap_mask: Colormap for masks (default 'gray')\n voxel_size: Voxel spacing for aspect ratio. If None, aspect=1.\n\n Example::\n \n # After patch_inference()\n show_segmentation_comparison(\n image=val_img,\n ground_truth=val_gt,\n prediction=predictions[0],\n metric_value=results_df.iloc[0]['DSC'],\n anatomical_plane=2\n )\n \"\"\"\n validate_anatomical_plane(anatomical_plane)\n \n img_data = image.data if hasattr(image, 'data') else image\n gt_data = ground_truth.data if hasattr(ground_truth, 'data') else ground_truth\n pred_data = prediction.data if hasattr(prediction, 'data') else prediction\n \n if hasattr(img_data, 'cpu'): img_data = img_data.cpu()\n if hasattr(gt_data, 'cpu'): gt_data = gt_data.cpu()\n if hasattr(pred_data, 'cpu'): pred_data = pred_data.cpu()\n \n if slice_index is None:\n gt_np = gt_data[channel].numpy()\n slice_index = find_max_slice(gt_np, anatomical_plane)\n \n fig, axes = plt.subplots(1, 3, figsize=figsize)\n \n img_slice, img_aspect = _get_slice(img_data, channel, slice_index, anatomical_plane, voxel_size)\n gt_slice, gt_aspect = _get_slice(gt_data, channel, slice_index, anatomical_plane, voxel_size)\n pred_slice, pred_aspect = _get_slice(pred_data, channel, slice_index, anatomical_plane, voxel_size)\n \n axes[0].imshow(img_slice, cmap=cmap_img, aspect=img_aspect)\n axes[0].set_title('Input Image')\n axes[0].axis('off')\n \n axes[1].imshow(gt_slice, cmap=cmap_mask, aspect=gt_aspect)\n axes[1].set_title('Ground Truth')\n axes[1].axis('off')\n \n pred_title = 'Prediction'\n if metric_value is not None:\n pred_title = f'Prediction ({metric_name}: {metric_value:.4f})'\n \n axes[2].imshow(pred_slice, cmap=cmap_mask, aspect=pred_aspect)\n axes[2].set_title(pred_title)\n axes[2].axis('off')\n \n plt.tight_layout()" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72d68ebd", + "metadata": {}, + "outputs": [], + "source": "#| export\ndef show_mask_overlay(image, mask, ctx=None, channel: int = 0, slice_index: int = None,\n anatomical_plane: int = 2, voxel_size=None, alpha: float = 0.5,\n cmap_img: str = 'gray', cmap_mask: str = 'autumn',\n ax=None, figsize=None, title=None, **kwargs):\n \"\"\"Overlay `mask` on `image` on a single axis, using the same slicing and orientation as `.show()`.\n\n Draws the image, then the mask on top with zero voxels transparent, so it lines up with\n `MedImage.show()` / `MedMask.show()`. Reuses the same `_get_slice` those methods use.\n\n Args:\n image: Background image (MedImage or object with `.data` [C, H, W, D]).\n mask: Overlay mask (MedMask or object with `.data` [C, H, W, D]); zero voxels are transparent.\n ctx: Axis to draw on (fastai-style context); same role as `ax`.\n channel: Channel to display.\n slice_index: Slice to show; if None, uses `find_max_slice(mask, anatomical_plane)`.\n anatomical_plane: 0=sagittal, 1=coronal, 2=axial (default).\n voxel_size: Voxel spacing for aspect ratio; if None, aspect defaults to 1.\n alpha: Overlay opacity.\n cmap_img: Colormap for the background image.\n cmap_mask: Colormap for the mask overlay.\n ax: Axis to draw on (overrides `ctx`).\n figsize: Figure size when a new axis is created.\n title: Optional axis title.\n kwargs: Extra args forwarded to the background image `imshow`.\n\n Returns:\n The matplotlib axis with the overlay.\n \"\"\"\n validate_anatomical_plane(anatomical_plane)\n\n img_data = image.data if hasattr(image, 'data') else image\n mask_data = mask.data if hasattr(mask, 'data') else mask\n if hasattr(img_data, 'cpu'): img_data = img_data.cpu()\n if hasattr(mask_data, 'cpu'): mask_data = mask_data.cpu()\n\n if slice_index is None:\n slice_index = find_max_slice(mask_data[channel].numpy(), anatomical_plane)\n\n img_slice, aspect = _get_slice(img_data, channel, slice_index, anatomical_plane, voxel_size)\n mask_slice, _ = _get_slice(mask_data, channel, slice_index, anatomical_plane, voxel_size)\n\n ax = ax if ax is not None else ctx\n if ax is None:\n _, ax = plt.subplots(figsize=figsize)\n\n ax.imshow(img_slice, cmap=cmap_img, aspect=aspect, **kwargs)\n ax.imshow(np.ma.masked_where(mask_slice == 0, mask_slice), cmap=cmap_mask, alpha=alpha, aspect=aspect)\n\n if title is not None:\n ax.set_title(title)\n ax.axis('off')\n return ax" } ], "metadata": { diff --git a/research/vestibular_schwannoma/01_five_fold_cross_validation.ipynb b/research/vestibular_schwannoma/01_five_fold_cross_validation.ipynb index c7433a3..f4d4489 100644 --- a/research/vestibular_schwannoma/01_five_fold_cross_validation.ipynb +++ b/research/vestibular_schwannoma/01_five_fold_cross_validation.ipynb @@ -7,65 +7,21 @@ "source": [ "# Five-Fold Cross-Validation for Vestibular Schwannoma Segmentation\n", "\n", - "This notebook runs 5-fold cross-validation for three 3D segmentation models on a\n", - "vestibular schwannoma (VS) MRI segmentation task, using the patch-based workflow in\n", - "fastMONAI, built on fastai, MONAI, and TorchIO.\n", - "\n", - "## The task\n", - "\n", - "Vestibular Schwannoma is a benign tumor arising from the vestibulocochlear nerve that can\n", - "cause hearing loss and brainstem compression ([Kujawa et al., 2024](https://doi.org/10.3389/fncom.2024.1365727)).\n", - "For small- to medium-sized lesions, guidelines recommend either upfront radiosurgery or\n", - "observation until radiographic tumor growth is evident ([Dhayalan et al., 2023](https://doi.org/10.1001/jama.2023.12222)).\n", - "Accurate, reproducible delineation of the tumor on Contrast-Enhanced T1-weighted (CE-T1w)\n", - "MRI supports treatment planning and volumetric follow-up. We treat it as a binary 3D\n", - "segmentation problem: for every voxel, decide tumor (1) or background (0).\n", - "\n", - "## The dataset\n", - "\n", - "We use 346 CE-T1w cases pooled from the Queen Square cohort\n", - "([Shapey et al., 2021](https://doi.org/10.1038/s41597-021-01064-w)) and the public crossMoDA\n", - "challenge data ([Dorent et al., 2023](https://doi.org/10.1016/j.media.2022.102628)). Each\n", - "case has a paired ground-truth tumor mask. The tumors are small relative to the field of\n", - "view, which is why we train on patches sampled around the lesion rather than on whole\n", - "volumes.\n", - "\n", - "## What five-fold cross-validation buys us\n", - "\n", - "Cross-validation gives a more honest estimate of how a model generalizes than a single\n", - "train/validation split. The 346 cases are partitioned into five folds. We train five times\n", - "per model; each run holds out one fold for validation and trains on the other four.\n", - "Averaging the held-out scores across the five folds uses every case for validation exactly\n", - "once and reports a mean and a spread rather than a single lucky (or unlucky) number.\n", - "\n", - "The fold assignment is fixed ahead of time in the `fold` column of the dataset CSV\n", - "(values 1..5). Using a pre-assigned column, rather than reshuffling here, means every model\n", - "sees the exact same folds, so their scores are directly comparable, and the split is\n", - "reproducible across notebooks and machines. The CSV also carries a `split` column, but the\n", - "cross-validation intentionally ignores it: all 346 cases participate in CV.\n", - "\n", - "## The three models\n", - "\n", - "- **UNet** (MONAI built-in): the classic encoder-decoder with skip connections. A strong,\n", - " well-understood convolutional baseline.\n", - "- **DynUNet** (MONAI built-in): a dynamic, nnU-Net-style UNet\n", - " ([Isensee et al., 2024](https://doi.org/10.1007/978-3-031-72114-4_47)) with residual blocks\n", - " and deep supervision (auxiliary losses at several decoder resolutions), which tends to\n", - " train more stably on medical data.\n", - "- **SegMamba** (our fork): a state-space (Mamba) backbone for 3D segmentation. It replaces\n", - " the quadratic-cost attention of a transformer with a linear-time selective scan, aiming to\n", - " capture long-range 3D context efficiently.\n", - "\n", - "All three are trained through an identical data, loss, metric, and evaluation pipeline, so\n", - "any difference in scores reflects the model rather than the plumbing.\n", - "\n", - "Despite the high accuracy reported for automated VS segmentation, its real-world utility and\n", - "computational cost remain underexplored, and an apparent performance plateau raises the\n", - "question of whether newer architectures offer practical gains over a standard CNN\n", - "([Connor et al., 2025](https://doi.org/10.5152/iao.2025.241693); [Häußler et al., 2025](https://doi.org/10.1002/lary.31979)).\n", - "Comparing a standard convolutional baseline (UNet) against an nnU-Net-style model (DynUNet)\n", - "and a state-space model (SegMamba) on the same data is a controlled way to probe that\n", - "question. References cited throughout are listed at the end of this notebook." + "Trains and evaluates three 3D segmentation models (UNet, DynUNet, SegMamba) on vestibular schwannoma (VS) MRI with 5-fold cross-validation, using the fastMONAI patch workflow. VS is a benign tumor of the vestibulocochlear nerve; accurate delineation on Contrast-Enhanced T1-weighted (CE-T1w) MRI supports treatment planning and volumetric follow-up ([Kujawa et al., 2024](https://doi.org/10.3389/fncom.2024.1365727); [Dhayalan et al., 2023](https://doi.org/10.1001/jama.2023.12222)).\n", + "\n", + "## Task and data\n", + "\n", + "Binary segmentation: label every voxel tumor (1) or background (0). We use 346 CE-T1w cases pooled from the Queen Square cohort ([Shapey et al., 2021](https://doi.org/10.1038/s41597-021-01064-w)) and the crossMoDA challenge ([Dorent et al., 2023](https://doi.org/10.1016/j.media.2022.102628)), each with a ground-truth mask. Tumors are small relative to the field of view, so we train on patches sampled around the lesion.\n", + "\n", + "## Cross-validation and models\n", + "\n", + "The 346 cases are split into five folds via the `fold` column (1..5) of the dataset CSV; each run holds out one fold for validation and trains on the other four, so every case is validated once. A fixed column keeps the folds identical across models and reproducible; the `split` column is ignored.\n", + "\n", + "- **UNet** (MONAI): convolutional encoder-decoder baseline.\n", + "- **DynUNet** (MONAI): nnU-Net-style UNet ([Isensee et al., 2024](https://doi.org/10.1007/978-3-031-72114-4_47)) with residual blocks and deep supervision.\n", + "- **SegMamba** (our fork): a state-space backbone with linear-time selective scan for long-range 3D context.\n", + "\n", + "All three run through an identical data, loss, and evaluation pipeline, so score differences reflect the model, not the plumbing. Comparing a CNN baseline, an nnU-Net-style model, and a state-space model probes whether newer architectures beat a standard CNN on VS ([Connor et al., 2025](https://doi.org/10.5152/iao.2025.241693); [Häußler et al., 2025](https://doi.org/10.1002/lary.31979)). References are listed at the end." ] }, { @@ -73,12 +29,7 @@ "id": "65895c58", "metadata": {}, "source": [ - "## 1. Environment setup\n", - "\n", - "We import the fastMONAI patch API, MONAI networks and losses, and the exact metric\n", - "functions used during evaluation. Everything runs in a conda environment with fastMONAI and\n", - "its dependencies installed. We `chdir` into the project folder early so that the relative\n", - "image and mask paths in the CSV (`../nii_data/...`) resolve correctly." + "## 1. Environment setup" ] }, { @@ -87,7 +38,42 @@ "id": "af9ad685", "metadata": {}, "outputs": [], - "source": "import os\nimport gc\nimport json\nimport traceback\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchio as tio\n\nfrom monai.losses import DiceCELoss, DeepSupervisionLoss\nfrom monai.networks.layers import Norm\nfrom monai.networks.nets import UNet, DynUNet\n\nfrom fastMONAI.vision_all import *\n\n# SegMamba is an optional fork; UNet and DynUNet need none of it. Import it once here so a\n# missing fork is reported now, with the install command, instead of mid-sweep. (A fork that\n# imports but whose CUDA kernels are broken surfaces later, when the model is built, and is\n# caught by the driver's try/except.)\n_SEGMAMBA_INSTALL_HINT = (\n \" GPU (training): pip install 'segmamba-v2[gpu] @ git+https://github.com/skaliy/SegMamba-V2.git'\\n\"\n \" CPU (inference): pip install 'segmamba-v2[cpu] @ git+https://github.com/skaliy/SegMamba-V2.git'\"\n)\ntry:\n from models_segmamba.segmambav2 import SegMamba\n SEGMAMBA_AVAILABLE = True\nexcept ImportError as exc:\n SegMamba, SEGMAMBA_AVAILABLE = None, False\n print(f\"[segmamba] fork not available ({exc}); UNet and DynUNet still run. To enable it:\\n\"\n + _SEGMAMBA_INSTALL_HINT)\n\n# Resolve data paths against the repo root so they work wherever the kernel started\n# (VS Code launches at the workspace root; plain Jupyter, in the notebook folder).\nimport fastMONAI\nREPO_ROOT = Path(fastMONAI.__file__).resolve().parent.parent\nos.chdir(REPO_ROOT / \"research\" / \"vestibular_schwannoma\")" + "source": [ + "import os\n", + "import gc\n", + "import json\n", + "import traceback\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import torch\n", + "import torchio as tio\n", + "\n", + "from monai.losses import DiceCELoss, DeepSupervisionLoss\n", + "from monai.networks.layers import Norm\n", + "from monai.networks.nets import UNet, DynUNet\n", + "\n", + "from fastMONAI.vision_all import *\n", + "\n", + "# Optional fork; import here so a missing fork is reported now with the install hint, not mid-sweep.\n", + "_SEGMAMBA_INSTALL_HINT = (\n", + " \" GPU (training): pip install 'segmamba-v2[gpu] @ git+https://github.com/skaliy/SegMamba-V2.git'\\n\"\n", + " \" CPU (inference): pip install 'segmamba-v2[cpu] @ git+https://github.com/skaliy/SegMamba-V2.git'\"\n", + ")\n", + "try:\n", + " from models_segmamba.segmambav2 import SegMamba\n", + " SEGMAMBA_AVAILABLE = True\n", + "except ImportError as exc:\n", + " SegMamba, SEGMAMBA_AVAILABLE = None, False\n", + " print(f\"[segmamba] fork not available ({exc}); UNet and DynUNet still run. To enable it:\\n\"\n", + " + _SEGMAMBA_INSTALL_HINT)\n", + "\n", + "# Resolve paths against the repo root so they work from any launch dir.\n", + "import fastMONAI\n", + "REPO_ROOT = Path(fastMONAI.__file__).resolve().parent.parent\n", + "os.chdir(REPO_ROOT / \"research\" / \"vestibular_schwannoma\")" + ] }, { "cell_type": "markdown", @@ -96,20 +82,7 @@ "source": [ "## 2. Run configuration\n", "\n", - "The knobs below control the size of the run. A full paper run is\n", - "`3 models x 5 folds x 500 epochs`, which takes on the order of days on a single modern GPU.\n", - "\n", - "To smoke-test the pipeline end to end in a few minutes instead, shrink the sweep, for\n", - "example:\n", - "\n", - "- `MODELS_TO_RUN = [\"unet\"]`\n", - "- `FOLDS_TO_RUN = [1]`\n", - "- `EPOCHS = 2`\n", - "\n", - "The demo exercises every stage (preprocessing, patch loading, training, sliding-window\n", - "evaluation, metric computation, aggregation) without producing publication-quality models.\n", - "`TARGET_SPACING` and `PATCH_SIZE` are part of the shared preprocessing contract and must\n", - "match the inference notebook, so they are not casual knobs." + "A full run is `3 models x 5 folds x 500 epochs`, on the order of days on one modern GPU. Shrink the knobs (e.g. `MODELS_TO_RUN = [\"unet\"]`, `FOLDS_TO_RUN = [1]`, `EPOCHS = 2`) for a quick smoke test that still exercises every stage. `TARGET_SPACING` and `PATCH_SIZE` are the shared preprocessing contract and must match `02_inference_new_cases.ipynb`." ] }, { @@ -119,25 +92,21 @@ "metadata": {}, "outputs": [], "source": [ - "# =============================== RUN KNOBS ===============================\n", - "# Which models and folds to run, and for how long. Shrink these for a quick\n", - "# pipeline smoke test (see the note above); keep the defaults for a full run.\n", "MODELS_TO_RUN = [\"unet\", \"dynunet\", \"segmamba\"]\n", "FOLDS_TO_RUN = [1, 2, 3, 4, 5]\n", "\n", - "EPOCHS = 500 # paper setting; set to e.g. 2 for a quick pipeline demo\n", - "BS = 4 # patches per training batch\n", - "LR = 1e-3 # maximum learning rate for fit_one_cycle\n", - "USE_TTA = True # 8-flip test-time augmentation during evaluation\n", - "# ========================================================================\n", + "EPOCHS = 500 # paper setting; 2 for a quick demo\n", + "BS = 4\n", + "LR = 1e-3 # max LR for fit_one_cycle\n", + "USE_TTA = True # 8-flip TTA at evaluation\n", "\n", - "# Shared preprocessing contract. These MUST match 02_inference_new_cases.ipynb.\n", + "# Shared preprocessing contract. MUST match 02_inference_new_cases.ipynb.\n", "TARGET_SPACING = [0.4102, 0.4102, 1.5] # resample target, voxel size in mm\n", "PATCH_SIZE = [192, 192, 48] # patch shape in voxels\n", "\n", "DATA_CSV = \"ml_dataset.csv\"\n", "\n", - "# cuDNN autotuner: a fixed patch size means the fastest kernels are found once and reused.\n", + "# cuDNN autotuner: fixed patch size means the fastest kernels are found once and reused.\n", "torch.backends.cudnn.benchmark = True\n", "\n", "print(f\"Models: {MODELS_TO_RUN}\")\n", @@ -153,11 +122,7 @@ "source": [ "## 3. Dataset and folds\n", "\n", - "The dataset CSV lists one row per case, with the raw image and mask paths, the pre-assigned\n", - "`fold` (1..5), and a `split` column. Cross-validation defines the validation set for each\n", - "fold inside `train_one_fold` as `is_val = (fold == fold_num)`; the held-out fold is the\n", - "validation set and the other four folds are training. All 346 cases take part in every\n", - "sweep, so the `split` column is deliberately not used here." + "The CSV lists one row per case with raw image and mask paths, the pre-assigned `fold` (1..5), and a `split` column. `train_one_fold` sets the validation set as `is_val = (fold == fold_num)`; all 346 cases take part, so the `split` column is not used here." ] }, { @@ -183,7 +148,11 @@ "cell_type": "markdown", "id": "e3e3534f", "metadata": {}, - "source": "## 4. Preprocess once to disk\n\nPreprocessing (RAS+ reorientation, resampling to `TARGET_SPACING`, and foreground-masked\nZ-normalization) is **fold-independent**: it depends only on each raw volume, not on which\nfold that volume lands in. The original per-model scripts repeated it inside every fold; we\nhoist it out and run it a single time over all 346 cases. This is faithful to the pipeline\nbut avoids four redundant passes per model.\n\n`preprocess_dataset` writes the processed volumes under `preprocessed/` and adds the columns\n`t1_img_path_preprocessed` and `t1_seg_path_preprocessed` to `train_df` in place. With\n`skip_existing=True` it is effectively idempotent: re-running reuses the cache instead of\nrecomputing, so you can safely point it at an existing `preprocessed/` folder. During\ntraining, `PatchConfig(preprocessed=True)` then skips reorder, resample, and normalization\nbecause the inputs are already prepared." + "source": [ + "## 4. Preprocess once to disk\n", + "\n", + "Preprocessing (RAS+ reorientation, resampling to `TARGET_SPACING`, foreground-masked Z-normalization) is fold-independent, so we run it once over all 346 cases. `preprocess_dataset` writes processed volumes under `preprocessed/` and adds `t1_img_path_preprocessed` / `t1_seg_path_preprocessed` to `train_df` in place; `skip_existing=True` makes it idempotent, reusing the cache on re-run. During training, `PatchConfig(preprocessed=True)` then skips reorder, resample, and normalization." + ] }, { "cell_type": "markdown", @@ -192,11 +161,7 @@ "source": [ "### Intensity normalization\n", "\n", - "`ZNormalization(masking_method=\"foreground\")` applies a Z-score: it subtracts a mean and\n", - "divides by a standard deviation. These are computed from foreground voxels, meaning voxels\n", - "with intensity above zero (read from the image, not the tumor mask), and then applied to the\n", - "whole volume. Restricting the statistics to the foreground stops the large zero-valued\n", - "background from dominating them and washing out tissue contrast." + "`ZNormalization(masking_method=\"foreground\")` computes mean and std from foreground voxels (intensity above zero, read from the image not the mask) and applies them to the whole volume. Restricting the statistics to the foreground stops the large zero-valued background from washing out tissue contrast." ] }, { @@ -206,9 +171,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Single source of truth for pre-patch / pre-inference intensity normalization. The same\n", - "# list is reused when building the training DataLoaders and when running sliding-window\n", - "# inference, guaranteeing identical preprocessing on both sides.\n", + "# Single source of truth for pre-patch / pre-inference intensity normalization.\n", "pre_patch_tfms = [ZNormalization(masking_method=\"foreground\")]\n", "\n", "preprocess_dataset(\n", @@ -229,7 +192,16 @@ "cell_type": "markdown", "id": "23c25507", "metadata": {}, - "source": "## 5. Shared pipeline building blocks\n\nEvery model is trained through the same two helpers, so differences in results come from the\nnetwork, not the data pipeline.\n\n`create_patch_config()` returns a `PatchConfig`, the single source of truth for how patches\nare sampled and stitched back together:\n\n- The volumes are large and the tumor is tiny, so we train on **192 x 192 x 48** patches\n rather than whole scans.\n- A **label sampler** with `label_probabilities={0: 0.2, 1: 0.8}` draws 80% of patches\n centered on tumor voxels. Without this bias the network would rarely see foreground.\n- `samples_per_volume=4`: four patches are drawn each time a volume is visited.\n- At evaluation time patches overlap by 50% and are blended with a **Hann window**\n (`aggregation_mode=\"hann\"`) for seamless reconstruction.\n- `preprocessed=True` tells the loader the volumes on disk are already reoriented, resampled,\n and normalized, so it does not repeat that work.\n- `normalization=pre_patch_tfms` records the intensity normalization (foreground\n Z-normalization) on the config as the single source of truth. The volumes on disk are\n already normalized, so training does not re-apply it under `preprocessed=True`; keeping it\n on the config documents the choice, logs it to the MLflow run, and lets inference rebuild the\n identical transform.\n- `keep_largest_component=True` post-processes predictions to the single largest connected\n region, which suits a solitary VS tumor." + "source": [ + "## 5. Shared pipeline building blocks\n", + "\n", + "Every model trains through the same helpers, so result differences come from the network, not the pipeline. `create_patch_config()` returns the `PatchConfig` that governs how patches are sampled and stitched:\n", + "\n", + "- A **label sampler** with `label_probabilities={0: 0.2, 1: 0.8}` draws 80% of patches centered on tumor voxels; without this bias the network would rarely see foreground.\n", + "- `preprocessed=True` tells the loader the volumes on disk are already reoriented, resampled, and normalized, so it does not repeat that work.\n", + "- `normalization=pre_patch_tfms` records the intensity normalization on the config as the single source of truth: not re-applied under `preprocessed=True`, but logged to MLflow and used to rebuild the identical transform at inference.\n", + "- `keep_largest_component=True` post-processes each prediction to its single largest connected region, which suits a solitary VS tumor." + ] }, { "cell_type": "code", @@ -237,7 +209,46 @@ "id": "42d6caca", "metadata": {}, "outputs": [], - "source": "def create_patch_config():\n \"\"\"Patch sampling and aggregation settings shared by all three models.\"\"\"\n return PatchConfig(\n patch_size=PATCH_SIZE,\n samples_per_volume=4,\n sampler_type=\"label\",\n label_probabilities={0: 0.2, 1: 0.8},\n patch_overlap=0.5,\n keep_largest_component=True,\n target_spacing=TARGET_SPACING,\n preprocessed=True,\n normalization=pre_patch_tfms,\n aggregation_mode=\"hann\",\n queue_num_workers=16,\n queue_length=1200,\n )\n\n\npatch_config = create_patch_config()\npatch_config" + "source": [ + "def create_patch_config():\n", + " \"\"\"Patch sampling and aggregation settings shared by all three models.\"\"\"\n", + " return PatchConfig(\n", + " patch_size=PATCH_SIZE,\n", + " samples_per_volume=4,\n", + " sampler_type=\"label\",\n", + " label_probabilities={0: 0.2, 1: 0.8},\n", + " patch_overlap=0.5,\n", + " keep_largest_component=True,\n", + " target_spacing=TARGET_SPACING,\n", + " preprocessed=True,\n", + " normalization=pre_patch_tfms,\n", + " aggregation_mode=\"hann\",\n", + " queue_num_workers=16,\n", + " queue_length=1200,\n", + " )\n", + "\n", + "\n", + "patch_config = create_patch_config()\n", + "\n", + "# Persist the config so 02_inference_new_cases.ipynb rebuilds identical preprocessing and patch\n", + "# settings. Local copy here; train_one_fold also logs it as a per-run MLflow artifact.\n", + "store_patch_variables(\n", + " \"inference_patch_config.json\",\n", + " patch_size=patch_config.patch_size,\n", + " patch_overlap=patch_config.patch_overlap,\n", + " aggregation_mode=patch_config.aggregation_mode,\n", + " apply_reorder=patch_config.apply_reorder,\n", + " target_spacing=patch_config.target_spacing,\n", + " sampler_type=patch_config.sampler_type,\n", + " label_probabilities=patch_config.label_probabilities,\n", + " samples_per_volume=patch_config.samples_per_volume,\n", + " queue_length=patch_config.queue_length,\n", + " queue_num_workers=patch_config.queue_num_workers,\n", + " keep_largest_component=patch_config.keep_largest_component,\n", + " normalization=patch_config.normalization,\n", + ")\n", + "patch_config" + ] }, { "cell_type": "markdown", @@ -246,12 +257,7 @@ "source": [ "### GPU augmentation\n", "\n", - "`create_gpu_augmentation()` returns a `GpuPatchAugmentation` that applies spatial and\n", - "intensity augmentation to each batch directly on the GPU, which keeps the input pipeline\n", - "from becoming the bottleneck. The transforms and probabilities follow nnU-Net conventions:\n", - "moderate random affine, anisotropy, flips, gamma, intensity scaling, noise, and blur.\n", - "Affine translations are specified in voxels, so the requested millimeter shifts are divided\n", - "by the target spacing. Augmentation is applied to the training set only." + "`create_gpu_augmentation()` returns a `GpuPatchAugmentation` that augments each batch on the GPU so the input pipeline is not the bottleneck. Transforms and probabilities follow nnU-Net conventions; affine translations are given in voxels, so millimeter shifts are divided by the target spacing, and augmentation is applied to the training set only." ] }, { @@ -281,11 +287,69 @@ " )" ] }, + { + "cell_type": "markdown", + "id": "525b796f", + "metadata": {}, + "source": [ + "### Sanity check: inspect a training batch\n", + "\n", + "Draw one batch of patches from the first fold and overlay the ground-truth mask, a quick check that the label sampler centers patches on the tumor, the intensities look normalized, and the mask lines up with the lesion. The cell is standalone and does not touch `train_one_fold` or the driver." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8f22431b", + "metadata": {}, + "outputs": [], + "source": [ + "# Pre-flight sanity check (standalone, runs once).\n", + "sanity_fold = FOLDS_TO_RUN[0]\n", + "sanity_df = train_df.copy()\n", + "sanity_df[\"is_val\"] = sanity_df[\"fold\"] == sanity_fold\n", + "\n", + "\n", + "sanity_config = create_patch_config() # fresh config leaves patch_config alone\n", + "sanity_config.samples_per_volume = 1\n", + "sanity_config.queue_length = BS\n", + "sanity_config.queue_num_workers = 2\n", + "\n", + "# No gpu_augmentation: show clean patches. preprocessed=True means they carry the on-disk\n", + "# foreground Z-norm, exactly what the model trains on.\n", + "sanity_dls = MedPatchDataLoaders.from_df(\n", + " df=sanity_df,\n", + " img_col=\"t1_img_path_preprocessed\",\n", + " mask_col=\"t1_seg_path_preprocessed\",\n", + " valid_col=\"is_val\",\n", + " patch_config=sanity_config,\n", + " bs=BS,\n", + ")\n", + "\n", + "# Axial slice per patch, mask overlaid. Look for patches centered on the lesion, near\n", + "# zero-mean contrast, and the mask hugging the tumor.\n", + "sanity_dls.show_batch(dl_idx=0, max_n=4, overlay=True, anatomical_plane=2)\n", + "\n", + "sanity_dls.close()" + ] + }, { "cell_type": "markdown", "id": "6e13dc69", "metadata": {}, - "source": "## 6. Model definitions\n\nThe three models are registered in a single `MODELS` dictionary. Each entry provides a\n`make_model` factory, a `make_loss` factory, and the exact experiment name, checkpoint\nfilename, and results directory that the sibling inference notebook expects:\n\n| key | experiment | checkpoint | results dir |\n|-----|------------|-----------|-------------|\n| `unet` | `vs5f_unet` | `best_unet` | `cv_results_unet` |\n| `dynunet` | `vs5f_dynunet` | `best_dynunet` | `cv_results_dynunet` |\n| `segmamba` | `vs5f_segmamba` | `best_segmamba` | `cv_results_segmamba` |\n\nUNet and SegMamba use a Dice + cross-entropy loss. DynUNet emits predictions at several\ndecoder depths, so it is wrapped in `DynUNetDSAdapter` (which turns the stacked output into\nthe list that `DeepSupervisionLoss` expects) and trained with a deep-supervision loss. UNet\nand DynUNet are compiled with `torch.compile`; SegMamba is not, because it relies on custom\nCUDA kernels.\n\n### Install SegMamba\n\n**UNet and DynUNet need nothing extra** beyond fastMONAI: they are MONAI built-ins.\n\n**SegMamba requires our fork** (`models_segmamba`, imported as `from\nmodels_segmamba.segmambav2 import SegMamba`; default backend `mamba_ssm`, 138.77M\nparameters). Install it from GitHub with the `[gpu]` extra, which is the one command you\nrun for this notebook: it brings the base package plus the CUDA Mamba kernels that the\ndefault `mamba_ssm` backend needs:\n\n pip install \"segmamba-v2[gpu] @ git+https://github.com/skaliy/SegMamba-V2.git\"\n\nThe kernels (`mamba-ssm`, `causal-conv1d`) sit behind the `[gpu]` extra rather than in the\nbase install because they compile CUDA extensions and only build where the CUDA toolchain is\npresent; keeping them optional lets the package install on machines without a GPU. For a\nCPU-only box (inference/dev; CPU training is not practical) use `[cpu]` instead, which pulls\nin transformers for the `mamba_backend=\"mambamixer\"` path:\n\n pip install \"segmamba-v2[cpu] @ git+https://github.com/skaliy/SegMamba-V2.git\"\n\nCaveat: `mamba-ssm` and `causal-conv1d` are CUDA-only and are compiled against the active\ntorch build. A torch version bump can break `causal-conv1d` with an \"undefined symbol\" ABI\nerror; reinstall it against the current torch with `--no-build-isolation --force-reinstall`.\n\nSegMamba is imported once, up front, in the setup cell (section 1) inside a `try/except`. If\nthe fork is not importable, `SEGMAMBA_AVAILABLE` is set to `False`, the install command is\nprinted, and the next cell skips SegMamba while UNet and DynUNet still run. A fork that imports\nbut whose CUDA kernels are broken (the \"undefined symbol\" case above) is not caught here,\nbecause the import does not load the kernels; it surfaces later, when the model is built, where\nthe driver's per-(model, fold) `try/except` skips just that run." + "source": [ + "## 6. Model definitions\n", + "\n", + "The three models are registered in a single `MODELS` dictionary. Each entry provides a `make_model` factory, a `make_loss` factory, and the experiment name, checkpoint filename, and results directory the inference notebook expects:\n", + "\n", + "| key | experiment | checkpoint | results dir |\n", + "|-----|------------|-----------|-------------|\n", + "| `unet` | `vs5f_unet` | `best_unet` | `cv_results_unet` |\n", + "| `dynunet` | `vs5f_dynunet` | `best_dynunet` | `cv_results_dynunet` |\n", + "| `segmamba` | `vs5f_segmamba` | `best_segmamba` | `cv_results_segmamba` |\n", + "\n", + "UNet and SegMamba use a Dice + cross-entropy loss. DynUNet emits predictions at several decoder depths, so it is wrapped in `DynUNetDSAdapter` and trained with `DeepSupervisionLoss`. UNet and DynUNet are compiled with `torch.compile`; SegMamba is not, because it relies on custom CUDA kernels." + ] }, { "cell_type": "code", @@ -350,8 +414,7 @@ " experiment=\"vs5f_dynunet\", best_fname=\"best_dynunet\", results_dir=\"cv_results_dynunet\",\n", ")\n", "\n", - "# SegMamba (optional fork): registered only if it imported cleanly in the setup cell. Not\n", - "# wrapped in torch.compile because it relies on custom CUDA kernels.\n", + "# SegMamba (optional fork): registered only if it imported cleanly. No torch.compile (custom CUDA kernels).\n", "if SEGMAMBA_AVAILABLE:\n", "\n", " def make_segmamba():\n", @@ -376,8 +439,7 @@ "\n", "print(\"Registered models:\", list(MODELS))\n", "\n", - "# If segmamba was requested but is unavailable, note it (the setup cell already printed the\n", - "# install/fix command). A UNet/DynUNet-only run stays quiet.\n", + "# Note a requested-but-unavailable segmamba (setup cell already printed the install hint).\n", "if \"segmamba\" in MODELS_TO_RUN and \"segmamba\" not in MODELS:\n", " print(\"\\n[segmamba] requested but its fork is not importable, so it is skipped \"\n", " \"(see the install hint from the setup cell above).\")" @@ -390,22 +452,14 @@ "source": [ "## 7. Train one fold\n", "\n", - "`train_one_fold(model_key, fold_num, train_df, patch_config)` runs a single training job. It\n", - "marks the held-out fold as validation (`is_val = fold == fold_num`), builds the patch\n", - "DataLoaders from the **preprocessed** columns, constructs the model and loss from the\n", - "registry, and trains with `fit_one_cycle`. Training uses:\n", - "\n", - "- `AccumulatedDice(n_classes=2)` as the monitored metric (nnU-Net-style pseudo-Dice\n", - " accumulated over the whole validation set, more stable than per-batch Dice).\n", - "- `EMACheckpoint` to save the best weights by the exponential moving average of that metric,\n", - " under the per-model checkpoint name so models do not overwrite each other.\n", - "- `create_mlflow_callback` to log parameters, metrics, the split, and artifacts to the\n", - " per-model MLflow experiment, tagged with the dataset fingerprint.\n", + "`train_one_fold(model_key, fold_num, train_df, patch_config)` runs one training job: it marks the held-out fold as validation, builds the patch DataLoaders from the preprocessed columns, constructs the model and loss from the registry, and trains with `fit_one_cycle`. Training uses:\n", + "\n", + "- `AccumulatedDice(n_classes=2)` as the monitored metric (nnU-Net-style pseudo-Dice accumulated over the whole validation set).\n", + "- `EMACheckpoint` to save the best weights by the EMA of that metric, under the per-model checkpoint name.\n", + "- `create_mlflow_callback` to log parameters, metrics, the split, and artifacts to the per-model MLflow experiment.\n", "- `.to_bf16()` for bfloat16 mixed-precision training.\n", "\n", - "At the end it calls `evaluate_fold` (defined in the next section) and frees GPU memory. The\n", - "function is defined here but only executed later by the driver loop, so it can safely refer\n", - "to `evaluate_fold` from the following cell." + "At the end it calls `evaluate_fold` (next section) and frees GPU memory." ] }, { @@ -414,7 +468,80 @@ "id": "5384dd04", "metadata": {}, "outputs": [], - "source": "def train_one_fold(model_key, fold_num, train_df, patch_config):\n \"\"\"Train one model on one fold, then evaluate the held-out fold.\"\"\"\n entry = MODELS[model_key]\n print(f\"\\n{'='*60}\")\n print(f\" {model_key.upper()} | FOLD {fold_num}\")\n print(f\"{'='*60}\\n\")\n\n # Fold membership: the held-out fold becomes the validation set. All 346 cases\n # participate; the CSV 'split' column is not used here.\n train_df = train_df.copy()\n train_df[\"is_val\"] = train_df[\"fold\"] == fold_num\n\n # Content fingerprint of the label set, logged to MLflow for reproducibility.\n med_dataset = MedDataset(\n img_list=train_df.t1_seg_path.tolist(), dtype=MedMask, max_workers=32\n )\n\n gpu_aug = create_gpu_augmentation()\n\n # Normalization lives on patch_config (single source of truth). The preprocessed columns\n # are already reoriented, resampled, and Z-normalized, so preprocessed=True skips re-applying\n # it here; the config still carries the spec for logging and inference parity.\n dls = MedPatchDataLoaders.from_df(\n df=train_df,\n img_col=\"t1_img_path_preprocessed\",\n mask_col=\"t1_seg_path_preprocessed\",\n valid_col=\"is_val\",\n patch_config=patch_config,\n gpu_augmentation=gpu_aug,\n bs=BS,\n )\n\n print(f\"[{model_key} fold {fold_num}] Train: {len(dls.train.subjects_dataset)}, \"\n f\"Val: {len(dls.valid.subjects_dataset)}\")\n\n model = entry[\"make_model\"]()\n loss_func = entry[\"make_loss\"]()\n\n learn = Learner(dls, model, loss_func=loss_func,\n metrics=[AccumulatedDice(n_classes=2)]).to_bf16()\n\n save_best = EMACheckpoint(\n monitor=\"accumulated_dice\", momentum=0.9,\n comp=np.greater, fname=entry[\"best_fname\"], with_opt=False,\n )\n\n mlflow_cb = create_mlflow_callback(\n learn,\n experiment_name=entry[\"experiment\"],\n run_name=f\"fold_{fold_num}\",\n extra_tags={\"fold\": str(fold_num)},\n dataset_version=med_dataset.fingerprint,\n )\n\n learn.fit_one_cycle(EPOCHS, LR, cbs=[mlflow_cb, save_best])\n\n results_df = evaluate_fold(\n learn, patch_config, dls, pre_patch_tfms, fold_num,\n tta=USE_TTA, mlflow_cb=mlflow_cb, results_dir=Path(entry[\"results_dir\"]),\n )\n\n # Release GPU memory before the next (model, fold).\n del learn, model, dls, gpu_aug\n torch.cuda.empty_cache()\n gc.collect()\n\n return results_df" + "source": [ + "def train_one_fold(model_key, fold_num, train_df, patch_config):\n", + " \"\"\"Train one model on one fold, then evaluate the held-out fold.\"\"\"\n", + " entry = MODELS[model_key]\n", + " print(f\"\\n{'='*60}\")\n", + " print(f\" {model_key.upper()} | FOLD {fold_num}\")\n", + " print(f\"{'='*60}\\n\")\n", + "\n", + " # Held-out fold becomes the validation set.\n", + " train_df = train_df.copy()\n", + " train_df[\"is_val\"] = train_df[\"fold\"] == fold_num\n", + "\n", + " # Content fingerprint of the label set, logged to MLflow for reproducibility.\n", + " med_dataset = MedDataset(\n", + " img_list=train_df.t1_seg_path.tolist(), dtype=MedMask, max_workers=32\n", + " )\n", + "\n", + " gpu_aug = create_gpu_augmentation()\n", + "\n", + " # Normalization lives on patch_config; preprocessed columns are already normalized, so\n", + " # preprocessed=True skips re-applying it here.\n", + " dls = MedPatchDataLoaders.from_df(\n", + " df=train_df,\n", + " img_col=\"t1_img_path_preprocessed\",\n", + " mask_col=\"t1_seg_path_preprocessed\",\n", + " valid_col=\"is_val\",\n", + " patch_config=patch_config,\n", + " gpu_augmentation=gpu_aug,\n", + " bs=BS,\n", + " )\n", + "\n", + " print(f\"[{model_key} fold {fold_num}] Train: {len(dls.train.subjects_dataset)}, \"\n", + " f\"Val: {len(dls.valid.subjects_dataset)}\")\n", + "\n", + " model = entry[\"make_model\"]()\n", + " loss_func = entry[\"make_loss\"]()\n", + "\n", + " learn = Learner(dls, model, loss_func=loss_func,\n", + " metrics=[AccumulatedDice(n_classes=2)]).to_bf16()\n", + "\n", + " save_best = EMACheckpoint(\n", + " monitor=\"accumulated_dice\", momentum=0.9,\n", + " comp=np.greater, fname=entry[\"best_fname\"], with_opt=False,\n", + " )\n", + "\n", + " mlflow_cb = create_mlflow_callback(\n", + " learn,\n", + " experiment_name=entry[\"experiment\"],\n", + " run_name=f\"fold_{fold_num}\",\n", + " extra_tags={\"fold\": str(fold_num)},\n", + " dataset_version=med_dataset.fingerprint,\n", + " )\n", + "\n", + " learn.fit_one_cycle(EPOCHS, LR, cbs=[mlflow_cb, save_best])\n", + "\n", + " # Reopen this fold's MLflow run to attach the patch config so 02 can pull it bound to these\n", + " # checkpoints (the callback closed the run but keeps its id).\n", + " import mlflow\n", + " if mlflow_cb.run_id is not None:\n", + " with mlflow.start_run(run_id=mlflow_cb.run_id):\n", + " mlflow.log_artifact(\"inference_patch_config.json\", artifact_path=\"config\")\n", + "\n", + " results_df = evaluate_fold(\n", + " learn, patch_config, dls, pre_patch_tfms, fold_num,\n", + " tta=USE_TTA, mlflow_cb=mlflow_cb, results_dir=Path(entry[\"results_dir\"]),\n", + " )\n", + "\n", + " # Release GPU memory before the next (model, fold).\n", + " del learn, model, dls, gpu_aug\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", + "\n", + " return results_df" + ] }, { "cell_type": "markdown", @@ -423,16 +550,7 @@ "source": [ "## 8. Evaluate a fold\n", "\n", - "After training, we score the held-out fold on the **raw** validation images (not the cached\n", - "preprocessed copies), re-applying the identical `ZNormalization` through\n", - "`pre_inference_tfms`. This mirrors real inference on unseen data end to end.\n", - "\n", - "For each case we compute Dice, sensitivity, precision, lesion-detection rate, and signed\n", - "relative volume error, plus spacing-aware surface metrics (ASSD, HD95, and Normalized Surface\n", - "Dice at 0.5/1.0/2.0 mm) via `calculate_surface_metrics`. Voxel spacing is read **per case**\n", - "from the ground-truth file, because the cohort spacing is not uniform and a single latched\n", - "affine would corrupt every distance. Per-case scores go to `results.csv`, provenance to\n", - "`benchmark_meta.json`, and finite-mean aggregates to MLflow." + "After training, we score the held-out fold on the raw validation images (not the preprocessed copies), re-applying the identical `ZNormalization` through `pre_inference_tfms`, which mirrors real inference end to end. For each case we compute Dice, sensitivity, precision, lesion-detection rate, and signed relative volume error, plus spacing-aware surface metrics (ASSD, HD95, and Normalized Surface Dice at 0.5/1.0/2.0 mm) via `calculate_surface_metrics`. Voxel spacing is read per case from the ground-truth file, because the cohort spacing is not uniform." ] }, { @@ -465,8 +583,7 @@ " tta=tta,\n", " )\n", "\n", - " # Per-case metric table via the library helper: it reads each GT mask's data AND spacing\n", - " # from one object (so the array and its spacing can't disagree) and runs the full panel.\n", + " # Per-case metric table: reads each GT mask's data and spacing from one object so they can't disagree.\n", " results_df = evaluate_segmentations(predictions, val_mask_paths,\n", " case_ids=val_df[\"case_id\"].tolist())\n", " results_df.insert(1, \"image\", [Path(p).name for p in val_img_paths])\n", @@ -475,8 +592,7 @@ " f\"({results_df['spacing_mm'].iloc[0]}); verify per-case derivation if unexpected.\")\n", " results_df.to_csv(fold_dir / \"results.csv\", index=False)\n", "\n", - " # Benchmark provenance (arxiv:2410.02630 transparency): metric implementation,\n", - " # NSD tolerances, and per-case spacing/empty-status used for the surface metrics.\n", + " # Benchmark provenance (arxiv:2410.02630 transparency): metric impl, NSD tolerances, per-case spacing.\n", " from fastMONAI.vision_metrics import _SURFACE_DISTANCE_SOURCE\n", " with open(fold_dir / \"benchmark_meta.json\", \"w\") as f:\n", " json.dump({\n", @@ -489,8 +605,7 @@ " for r in results_df.to_dict(\"records\")],\n", " }, f, indent=2)\n", "\n", - " # inf-safe means: one_empty cases score assd_mm/hd95_mm = inf, which would poison a plain\n", - " # mean (and MLflow rejects non-finite values). Replacing inf with NaN lets pandas skip it.\n", + " # inf-safe means: one_empty cases score inf; replace with NaN so pandas skips it (MLflow rejects non-finite).\n", " numeric = results_df.select_dtypes(include=\"number\").replace([np.inf, -np.inf], np.nan)\n", " mlflow_cb.log_metrics_table(results_df, display=False)\n", " mlflow_cb.log_metrics({f\"val_{m}\": numeric[m].mean() for m in numeric.columns})\n", @@ -508,14 +623,7 @@ "source": [ "## 9. Cross-validation driver\n", "\n", - "The driver sweeps `MODELS_TO_RUN x FOLDS_TO_RUN`, calling `train_one_fold` for each pair.\n", - "Each pair is wrapped in try/except: if one job fails (for example an out-of-memory error on\n", - "one fold) the error is printed and the sweep continues with the next pair rather than\n", - "aborting the whole run. Unregistered model keys (for example `segmamba` without the fork)\n", - "are skipped with a message.\n", - "\n", - "This is the long-running cell. For the full paper configuration it trains up to 15 models\n", - "and can take days. Reduce the knobs in the configuration cell for a quick demo." + "The driver sweeps `MODELS_TO_RUN x FOLDS_TO_RUN`, calling `train_one_fold` for each pair. Each pair is wrapped in try/except, so one failure (for example an out-of-memory error on one fold) is logged and the sweep continues rather than aborting." ] }, { @@ -525,8 +633,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Resilient sweep: a failure on one (model, fold) is logged and skipped so it does not abort\n", - "# the rest of the run. Re-running is safe; a completed fold simply overwrites its outputs.\n", + "# Resilient sweep: a failure on one (model, fold) is logged and skipped, not fatal.\n", "for model_key in MODELS_TO_RUN:\n", " if model_key not in MODELS:\n", " print(f\"[skip] '{model_key}' is not registered (see the model definitions cell above).\")\n", @@ -554,12 +661,7 @@ "source": [ "## 10. Aggregate per-model results\n", "\n", - "For each model, `aggregate_results` concatenates the per-fold `results.csv` files into a\n", - "single `cv_summary.csv` inside that model's results directory, then prints a per-metric\n", - "mean/standard-deviation summary and the per-fold DSC breakdown. Surface distances in\n", - "millimeters (ASSD, HD95) are averaged over **finite values only**: a case where exactly one\n", - "of prediction or ground truth is empty scores `+inf` by design, and a plain mean would be\n", - "dominated by it." + "For each model, `aggregate_results` concatenates the per-fold `results.csv` files into `cv_summary.csv` and prints a per-metric mean/std summary plus the per-fold DSC breakdown. Millimeter surface distances (ASSD, HD95) are averaged over finite values only, since a case where exactly one of prediction or ground truth is empty scores `+inf` by design." ] }, { @@ -570,11 +672,7 @@ "outputs": [], "source": [ "def aggregate_results(results_dir):\n", - " \"\"\"Concatenate a model's per-fold results.csv into cv_summary.csv and print a summary.\n", - "\n", - " Returns the combined DataFrame, or None if no fold results exist yet (notebook-safe:\n", - " it does not exit the process).\n", - " \"\"\"\n", + " \"\"\"Concatenate a model's per-fold results.csv into cv_summary.csv and print a summary; None if no results yet.\"\"\"\n", " results_dir = Path(results_dir)\n", " all_results = []\n", " for fold_dir in sorted(results_dir.glob(\"fold_*\")):\n", @@ -653,12 +751,7 @@ "source": [ "## 11. Cross-model comparison\n", "\n", - "This is the headline table and it is net-new relative to the per-model scripts. For every\n", - "model that produced results, we pool all held-out cases across folds and report the mean and\n", - "standard deviation of each key metric, save the numbers to `cv_model_comparison.csv`, and\n", - "display a compact `mean +/- std` view. The same inf-safe averaging is applied here, so the\n", - "millimeter surface metrics ignore non-finite per-case scores (and NaN, e.g. RVE on an empty\n", - "ground truth, is skipped the same way)." + "For every model that produced results, we pool all held-out cases across folds, report the mean and standard deviation of each key metric, save the numbers to `cv_model_comparison.csv`, and display a compact `mean +/- std` view. The same inf-safe averaging is applied, so the millimeter surface metrics ignore non-finite per-case scores (and NaN, e.g. RVE on an empty ground truth, is skipped the same way)." ] }, { @@ -668,10 +761,6 @@ "metadata": {}, "outputs": [], "source": [ - "# Headline metrics for the cross-model table. Surface distances in mm (assd/hd95) can be\n", - "# +inf for 'one_empty' cases (exactly one of prediction/ground truth empty). pandas .mean()/\n", - "# .std() already skip NaN but not inf, so we replace inf with NaN first; the finite-only mean\n", - "# and sample std (ddof=1, the pandas default) then match the per-model summary above.\n", "KEY_METRICS = [\"dsc\", \"sensitivity\", \"precision\", \"ldr\", \"rve\",\n", " \"assd_mm\", \"hd95_mm\", \"nsd_tau1.0_mm\"]\n", "\n", @@ -704,38 +793,127 @@ " print(\"No per-model results available yet. Run the driver loop first.\")" ] }, + { + "cell_type": "markdown", + "id": "cf1a1c7a", + "metadata": {}, + "source": [ + "## 12. Qualitative check: prediction vs. ground truth\n", + "\n", + "A visual sanity check on the saved validation masks. We reload one model's predictions from disk and compare the median-DSC (representative) and worst-DSC (failure) cases against their ground truth; input, ground truth, and prediction share each case's original voxel grid, so they line up slice for slice." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c39cd8cc", + "metadata": {}, + "outputs": [], + "source": [ + "from fastMONAI.vision_plot import *\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "def _pred_filename(input_path):\n", + " \"\"\"Prediction filename patch_inference wrote: '_pred.nii[.gz]'. Mirrors vision_patch._save_prediction.\"\"\"\n", + " p = Path(input_path)\n", + " if p.suffix == \".gz\" and p.stem.endswith(\".nii\"):\n", + " return f\"{p.stem[:-4]}_pred.nii.gz\"\n", + " if p.suffix == \".nii\":\n", + " return f\"{p.stem}_pred.nii\"\n", + " return f\"{p.stem}_pred.nii.gz\"\n", + "\n", + "\n", + "def show_qualitative_examples(plane=2):\n", + " \"\"\"Show the median-DSC and worst-DSC validation cases for the first model with results on disk.\"\"\"\n", + " # First registered model that actually produced fold results.\n", + " picked = None\n", + " for model_key, entry in MODELS.items():\n", + " results_dir = Path(entry[\"results_dir\"])\n", + " fold_dirs = [d for d in sorted(results_dir.glob(\"fold_*\"))\n", + " if (d / \"results.csv\").exists()]\n", + " if fold_dirs:\n", + " picked = (model_key, results_dir, fold_dirs)\n", + " break\n", + " if picked is None:\n", + " print(\"No fold_*/results.csv found for any model. Run the sweep first.\")\n", + " return\n", + " model_key, results_dir, fold_dirs = picked\n", + "\n", + " # Pool the per-fold tables, tagging each row with its validation fold (the predictions/ subdir).\n", + " frames = []\n", + " for fold_dir in fold_dirs:\n", + " df = pd.read_csv(fold_dir / \"results.csv\")\n", + " df[\"fold\"] = int(fold_dir.name.split(\"_\")[1])\n", + " frames.append(df)\n", + " pooled = pd.concat(frames, ignore_index=True)\n", + " print(f\"Model '{model_key}': pooled {len(pooled)} cases from folds \"\n", + " f\"{sorted(pooled['fold'].unique().tolist())}.\")\n", + "\n", + " # Sort by dsc: middle row is the median case, first is the worst (real cases, no interpolation).\n", + " ordered = pooled.sort_values(\"dsc\").reset_index(drop=True)\n", + " picks = [(\"median DSC\", ordered.iloc[len(ordered) // 2]),\n", + " (\"worst DSC\", ordered.iloc[0])]\n", + "\n", + " shown = set()\n", + " for label, row in picks:\n", + " case_id = row[\"case_id\"]\n", + " if case_id in shown: # median == worst when only one case\n", + " continue\n", + " shown.add(case_id)\n", + "\n", + " # case_id -> raw t1/seg paths (original grid) via train_df.\n", + " sub = train_df[train_df[\"case_id\"] == case_id]\n", + " if sub.empty:\n", + " print(f\"[skip] {case_id}: not found in train_df.\")\n", + " continue\n", + " img_path = str(sub.iloc[0][\"t1_img_path\"])\n", + " gt_path = str(sub.iloc[0][\"t1_seg_path\"])\n", + "\n", + " fold = int(row[\"fold\"])\n", + " pred_path = results_dir / f\"fold_{fold}\" / \"predictions\" / _pred_filename(img_path)\n", + " if not pred_path.exists():\n", + " print(f\"[skip] {case_id}: prediction not found at {pred_path}.\")\n", + " continue\n", + "\n", + " # Default loader (no reorder/resample): input, GT and prediction share the case's original\n", + " # grid. voxel_size is the native spacing, so the display aspect is correct.\n", + " img = MedImage.create(img_path)\n", + " gt = MedMask.create(gt_path)\n", + " pred = MedMask.create(str(pred_path))\n", + " vsize = tio.ScalarImage(img_path).spacing\n", + "\n", + " dsc = float(row[\"dsc\"])\n", + " print(f\"[{label}] {case_id} (fold {fold}): DSC = {dsc:.4f}\")\n", + " show_segmentation_comparison(img, gt, pred, metric_value=dsc, metric_name=\"DSC\",\n", + " anatomical_plane=plane, voxel_size=vsize)\n", + " plt.show()\n", + "\n", + "\n", + "show_qualitative_examples()" + ] + }, { "cell_type": "markdown", "id": "4b709ef9", "metadata": {}, "source": [ - "## 12. Viewing runs and next steps\n", + "## 13. Viewing runs and next steps\n", "\n", - "Every fold logs to MLflow (metrics, parameters, the train/validation split, and model\n", - "artifacts) under the per-model experiment names `vs5f_unet`, `vs5f_dynunet`, and\n", - "`vs5f_segmamba`. To browse them, launch the fastMONAI MLflow UI from a notebook cell:\n", + "Every fold logs metrics, parameters, the train/validation split, and model artifacts to MLflow under the per-model experiments `vs5f_unet`, `vs5f_dynunet`, and `vs5f_segmamba`. To browse them, launch the fastMONAI MLflow UI from a notebook cell:\n", "\n", " mlflow_ui = MLflowUIManager()\n", " mlflow_ui.start_ui() # opens http://localhost:5001\n", "\n", - "The UI is tied to the kernel and is reaped when the interpreter exits, so restarting the\n", - "notebook leaves no orphaned server holding the port.\n", - "\n", "Artifacts written to disk:\n", "\n", "- `preprocessed/` - reoriented, resampled, normalized volumes (shared by all folds).\n", - "- `cv_results_/fold_N/` - per-fold `results.csv`, `benchmark_meta.json`, and NIfTI\n", - " predictions.\n", + "- `inference_patch_config.json` - patch/preprocessing config, also logged to each fold's MLflow run under `config/`; notebook 02 loads it for inference parity.\n", + "- `cv_results_/fold_N/` - per-fold `results.csv`, `benchmark_meta.json`, and NIfTI predictions.\n", "- `cv_results_/cv_summary.csv` - all folds concatenated for one model.\n", "- `cv_model_comparison.csv` - the headline mean +/- std table across models.\n", "\n", - "The best checkpoint for each model is saved by `EMACheckpoint` as `best_unet`,\n", - "`best_dynunet`, and `best_segmamba`. To run any trained model on brand-new cases, see the\n", - "sibling notebook `02_inference_new_cases.ipynb`, which reuses the same `TARGET_SPACING`,\n", - "`PATCH_SIZE`, and `ZNormalization(masking_method=\"foreground\")` contract so that inference\n", - "preprocessing matches training exactly.\n", - "\n", - "To deploy the five folds together as a soft-vote ensemble, pass the list of fold learners to `patch_inference` (or `PatchInferenceEngine`): it averages their per-patch probabilities in a single sliding-window pass, which usually beats any single fold. Notebook 02 does exactly this." + "The best checkpoint per model is saved by `EMACheckpoint` as `best_unet`, `best_dynunet`, and `best_segmamba`. To run any trained model on new cases, see `02_inference_new_cases.ipynb`, which loads `inference_patch_config.json` so inference reuses the exact preprocessing contract. To deploy the five folds together as a soft-vote ensemble, pass the list of fold learners to `patch_inference`; notebook 02 does exactly this." ] }, { @@ -768,4 +946,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/research/vestibular_schwannoma/02_inference_new_cases.ipynb b/research/vestibular_schwannoma/02_inference_new_cases.ipynb index 3791831..967ed3f 100644 --- a/research/vestibular_schwannoma/02_inference_new_cases.ipynb +++ b/research/vestibular_schwannoma/02_inference_new_cases.ipynb @@ -4,18 +4,22 @@ "cell_type": "markdown", "id": "cfd2f9a8", "metadata": {}, - "source": "# Inference on New Vestibular Schwannoma Cases\n\nThis notebook segments new vestibular schwannoma (VS) MRI scans using the models trained in `01_five_fold_cross_validation.ipynb` and saves the predicted tumor masks as NIfTI files. The audience is students and researchers, so every step is explained.\n\nIt runs a **5-fold soft-voting ensemble**: all five cross-validation models for the chosen architecture are run on the scan, their class probabilities are averaged, and the averaged map is decoded into a mask. This is the standard way to deploy a cross-validation and usually beats any single fold.\n\n## What this notebook does\n\n1. Locates the five fold checkpoints for a model (UNet, DynUNet, or SegMamba) trained in notebook 01.\n2. Rebuilds the exact preprocessing used at training time.\n3. Runs sliding-window patch inference with each fold and averages the probabilities (soft voting).\n4. Saves and visualizes the predicted mask, and optionally scores it against ground truth.\n\n## Sliding-window patch inference\n\nVS scans are large 3D volumes, too big to feed to a 3D network in one piece. Instead we use **patch-based inference**:\n\n- A `GridSampler` slides a fixed-size window (the patch) across the volume with overlap, producing many small patches.\n- The network predicts on each patch.\n- A `GridAggregator` stitches the patch predictions back into a full-volume prediction. We use **Hann-window** aggregation, which weights each patch by a smooth cosine taper so overlapping patches blend without visible seams at the patch borders.\n\nfastMONAI wraps this in `PatchInferenceEngine` and the `patch_inference` helper.\n\n## Soft-voting ensemble\n\nA 5-fold cross-validation produces five models, each trained on a different 4/5 of the data. To turn them into one predictor we combine them by **soft voting**: run all five, apply softmax to each output, average the five probability maps, and only then take the argmax. Averaging probabilities (not the decoded masks) is the correct way to combine multi-class segmentation models, and it gives a small, reliable accuracy gain over any single fold. The engine returns each fold's probability map already resampled and reoriented into the input's original space, so the five maps are voxel-aligned and can be averaged directly.\n\n## The one rule that matters most: preprocessing parity\n\nA segmentation model produces correct output only when the data it sees at inference is preprocessed **identically** to the data it saw during training. Three settings must match exactly:\n\n- **`apply_reorder`**: reorientation to RAS+ canonical orientation.\n- **`target_spacing`**: the voxel spacing the volume is resampled to.\n- **`normalization`**: the intensity normalization (here, foreground Z-normalization).\n\nA mismatch does not raise an error. It silently produces wrong predictions, for example masks that are shifted, mirrored, or rotated. This notebook reuses the same `PatchConfig` and normalization as training so the two pipelines stay in lockstep." + "source": [ + "# Inference on New Vestibular Schwannoma Cases\n", + "\n", + "Segments new VS MRI scans with the models trained in `01_five_fold_cross_validation.ipynb` and saves the predicted tumor masks as NIfTI. It runs a **5-fold soft-voting ensemble**: all five fold models for the chosen architecture (UNet, DynUNet, or SegMamba) are run on the scan, their softmax probabilities averaged, and the average decoded into a mask (argmax + keep-largest-component). This usually beats any single fold.\n", + "\n", + "## Preprocessing parity\n", + "\n", + "A model gives correct output only when inference preprocesses the input exactly as training did. Three settings must match: `apply_reorder` (RAS+ reorientation), `target_spacing` (resample spacing), and `normalization` (foreground Z-normalization). A mismatch does not raise an error, it silently degrades the prediction." + ] }, { "cell_type": "markdown", "id": "106beda7", "metadata": {}, "source": [ - "## 1. Environment setup\n", - "\n", - "We import the fastMONAI public API and pin the working directory to the project folder so that the relative dataset paths used here (for example `../nii_data/...`) resolve the same way they do during training.\n", - "\n", - "`from fastMONAI.vision_all import *` brings in everything we need for inference: `PatchConfig`, `patch_inference`, `ZNormalization`, `store_patch_variables` / `load_patch_variables`, `MedImage`, `MedMask`, and the metric functions. `load_learner` comes from fastai and is used to load an exported model." + "## 1. Environment setup" ] }, { @@ -24,13 +28,30 @@ "id": "44c21482", "metadata": {}, "outputs": [], - "source": "import os\nfrom pathlib import Path\n\nimport numpy as np\nimport torch\nimport torchio as tio\n\nfrom fastMONAI.vision_all import *\nfrom fastai.learner import load_learner\n\n# Resolve data paths against the repo root so they work wherever the kernel started\n# (VS Code launches at the workspace root; plain Jupyter, in the notebook folder).\nimport fastMONAI\nREPO_ROOT = Path(fastMONAI.__file__).resolve().parent.parent\nos.chdir(REPO_ROOT / \"research\" / \"vestibular_schwannoma\")\n\nimport fastai, monai # fastMONAI already imported above\nprint(\"fastMONAI:\", fastMONAI.__version__)\nprint(\"fastai: \", fastai.__version__)\nprint(\"MONAI: \", monai.__version__)\nprint(\"TorchIO: \", tio.__version__)\nprint(\"PyTorch: \", torch.__version__)\nprint(\"CUDA available:\", torch.cuda.is_available())\nif torch.cuda.is_available():\n print(\"GPU:\", torch.cuda.get_device_name(0))" + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import torch\n", + "import torchio as tio\n", + "\n", + "from fastMONAI.vision_all import *\n", + "from fastai.learner import load_learner\n", + "\n", + "# Resolve paths against the repo root so they work from any launch dir.\n", + "import fastMONAI\n", + "REPO_ROOT = Path(fastMONAI.__file__).resolve().parent.parent\n", + "os.chdir(REPO_ROOT / \"research\" / \"vestibular_schwannoma\")" + ] }, { "cell_type": "markdown", "id": "a29cda31", "metadata": {}, - "source": "## 2. Configuration\n\nEverything you would normally change lives in the cell below. Treat this notebook as a **template**: pick the architecture with `MODEL_KEY`, list the scans in `NEW_CASES`, and run top to bottom. The notebook always soft-votes the five cross-validation folds for the chosen model.\n\n| Knob | Meaning |\n| --- | --- |\n| `MODEL_KEY` | Which architecture to load: `\"unet\"`, `\"dynunet\"`, or `\"segmamba\"`. Must match a trained CV experiment. |\n| `FOLD_LEARNER_PATHS` | Optional `{fold: path/to/best_learner.pkl}`. Leave empty to auto-discover the five folds from MLflow. |\n| `NEW_CASES` | List of NIfTI image paths to segment. These are **raw** scans; the pipeline preprocesses them for you. |\n| `OUTPUT_DIR` | Where predicted masks are written as NIfTI. |\n| `USE_TTA` | 8-flip test-time augmentation. More robust, roughly 8x slower. |\n| `USE_AMP` | Automatic mixed precision (float16) forward pass. Faster on CUDA, ignored on CPU. |\n\n`TARGET_SPACING` and `PATCH_SIZE` are part of the shared contract with training and must not be changed independently of notebook 01." + "source": [ + "## 2. Configuration" + ] }, { "cell_type": "code", @@ -38,13 +59,34 @@ "id": "5e2b4255", "metadata": {}, "outputs": [], - "source": "# --- Which model ---------------------------------------------------------------\nMODEL_KEY = \"unet\" # \"unet\" | \"dynunet\" | \"segmamba\" (must match a trained CV run)\n\n# Ensemble checkpoints: {fold_number: path/to/best_learner.pkl}. Leave empty to\n# auto-discover the five folds from the MLflow experiment \"vs5f_\" below.\nFOLD_LEARNER_PATHS = {}\n\n# --- What to segment and where to write it -------------------------------------\nNEW_CASES = [\n \"../nii_data/queen_square_data/vs_gk_1/vs_gk_1_t1_refT1.nii.gz\",\n]\nOUTPUT_DIR = \"inference_predictions\"\n\n# --- Inference options ---------------------------------------------------------\nUSE_TTA = True # 8-flip test-time augmentation (more robust, ~8x slower)\nUSE_AMP = True # mixed precision on CUDA (ignored on CPU)\n\n# --- Shared contract with training notebook 01 (do not change in isolation) ----\nTARGET_SPACING = [0.4102, 0.4102, 1.5]\nPATCH_SIZE = [192, 192, 48]\n\nprint(f\"Model: {MODEL_KEY} | mode: 5-fold soft-voting ensemble\")\nprint(f\"Cases to segment: {len(NEW_CASES)}\")" + "source": [ + "MODEL_KEY = \"unet\" # \"unet\" | \"dynunet\" | \"segmamba\"\n", + "\n", + "# leave empty to auto-discover from MLflow.\n", + "FOLD_LEARNER_PATHS = {}\n", + "\n", + "NEW_CASES = [\n", + " \"../nii_data/CASE_ID\", \n", + "]\n", + "OUTPUT_DIR = \"inference_predictions\"\n", + "\n", + "USE_TTA = True # 8-flip test-time augmentation (more robust, ~8x slower)\n", + "USE_AMP = True # mixed precision on CUDA (ignored on CPU)\n", + "\n", + "# Patch config: pulled from the model's MLflow run by default; set CONFIG_JSON to a local inference_patch_config.json to override.\n", + "CONFIG_JSON = None\n", + "\n", + "print(f\"Model: {MODEL_KEY} | mode: 5-fold soft-voting ensemble\")\n", + "print(f\"Cases to segment: {len(NEW_CASES)}\")" + ] }, { "cell_type": "markdown", "id": "e4287608", "metadata": {}, - "source": "### Locate the fold checkpoints\n\nNotebook 01 logs `best_learner.pkl` as an MLflow artifact for every fold, under experiments named `vs5f_unet`, `vs5f_dynunet`, and `vs5f_segmamba` (one run per fold, tagged with its fold number). The helper below finds one checkpoint per fold for `MODEL_KEY` and returns `{fold: local_path}`. It is best-effort and never raises, so it is safe to run before any training exists.\n\nIf you already have the paths, set `FOLD_LEARNER_PATHS` in the configuration cell and this lookup is skipped." + "source": [ + "### Locate the fold checkpoints" + ] }, { "cell_type": "code", @@ -52,25 +94,22 @@ "id": "712e4005", "metadata": {}, "outputs": [], - "source": "# find_fold_learners now lives in the library (fastMONAI.utils) and is imported via\n# `from fastMONAI.vision_all import *` above; the experiment name is passed explicitly\n# (it also respects a user-configured MLflow tracking server, unlike the old inline copy).\n\nif not FOLD_LEARNER_PATHS:\n FOLD_LEARNER_PATHS = find_fold_learners(f\"vs5f_{MODEL_KEY}\")\n\nprint(\"Fold checkpoints:\", {k: str(v) for k, v in sorted(FOLD_LEARNER_PATHS.items())} or \"(none)\")" + "source": [ + "# find_fold_learners (fastMONAI.utils) discovers one best_learner.pkl per fold from MLflow.\n", + "if not FOLD_LEARNER_PATHS:\n", + " FOLD_LEARNER_PATHS = find_fold_learners(f\"vs5f_{MODEL_KEY}\")\n", + "\n", + "print(\"Fold checkpoints:\", {k: str(v) for k, v in sorted(FOLD_LEARNER_PATHS.items())} or \"(none)\")" + ] }, { "cell_type": "markdown", "id": "b59bde21", "metadata": {}, "source": [ - "## 3. Build the inference configuration\n", - "\n", - "`PatchConfig` is the single object that describes how patches are sampled, how they are aggregated, and (critically) how the raw input is preprocessed. We build the **exact same** config that training used. The values below are the shared contract with notebook 01.\n", + "## 3. Load the inference configuration\n", "\n", - "A few points worth understanding:\n", - "\n", - "- **`preprocessed=True`** only affects *training* (it tells the training pipeline that the data on disk was already reordered, resampled, and normalized, so it should not do it again). At *inference* the engine always reorders, resamples, and normalizes the raw input for you, so we point it at the original scans.\n", - "- **`patch_overlap=0.5`** means neighboring windows overlap by half a patch. More overlap gives smoother, more accurate borders at the cost of speed.\n", - "- **`aggregation_mode=\"hann\"`** blends overlapping patches with a Hann taper (smoothest transitions).\n", - "- **`keep_largest_component=True`** keeps only the largest connected foreground blob, a simple and effective post-processing step for a single-tumor task like VS.\n", - "\n", - "The pre-inference normalization is defined separately as `pre_inference_tfms` and passed to `patch_inference`. It must be the same transform training used: foreground `ZNormalization`. Passing it explicitly makes the parity obvious and overrides whatever is stored on the config." + "`PatchConfig` describes how patches are sampled and aggregated and how the raw input is preprocessed. We load the exact config notebook 01 saved at training time (from the model's MLflow run by default, or a local `inference_patch_config.json`) and rebuild `PatchConfig` from it." ] }, { @@ -80,101 +119,79 @@ "metadata": {}, "outputs": [], "source": [ - "# Pre-inference intensity normalization. MUST match training exactly.\n", - "pre_inference_tfms = [ZNormalization(masking_method=\"foreground\")]\n", - "\n", - "# The same PatchConfig used for training in notebook 01.\n", - "patch_config = PatchConfig(\n", - " patch_size=PATCH_SIZE,\n", - " samples_per_volume=4,\n", - " sampler_type=\"label\",\n", - " label_probabilities={0: 0.2, 1: 0.8},\n", - " patch_overlap=0.5,\n", - " keep_largest_component=True,\n", - " target_spacing=TARGET_SPACING,\n", - " preprocessed=True, # only affects training; inference always preprocesses raw input\n", - " aggregation_mode=\"hann\",\n", - " queue_num_workers=16,\n", - " queue_length=1200,\n", - ")\n", + "# Load the patch config notebook 01 saved: the source of truth for preprocessing and patch sampling.\n", + "if CONFIG_JSON and Path(CONFIG_JSON).exists():\n", + " config_path = CONFIG_JSON\n", + " print(f\"Patch config: local file {config_path}\")\n", + "else:\n", + " _cfg = find_fold_learners(f\"vs5f_{MODEL_KEY}\",\n", + " artifact_path=\"config/inference_patch_config.json\")\n", + " if _cfg:\n", + " config_path = next(iter(_cfg.values())) # identical across folds; take any\n", + " print(f\"Patch config: MLflow run for 'vs5f_{MODEL_KEY}'\")\n", + " elif Path(\"inference_patch_config.json\").exists():\n", + " config_path = \"inference_patch_config.json\"\n", + " print(\"Patch config: local inference_patch_config.json (no MLflow artifact found)\")\n", + " else:\n", + " raise FileNotFoundError(\n", + " \"No patch config found. Run notebook 01 (it logs the config to MLflow and writes a \"\n", + " \"local copy), or set CONFIG_JSON to a local inference_patch_config.json.\")\n", + "\n", + "# normalization travels inside the config; patch_inference applies the same ZNormalization.\n", + "patch_config = PatchConfig(**load_patch_variables(config_path))\n", "print(patch_config)" ] }, { "cell_type": "markdown", - "id": "aec89045", + "id": "20f7b65d", "metadata": {}, "source": [ - "### Persisting and reloading the config\n", + "## 4. Load the models\n", "\n", - "To guarantee parity across machines and over time, the patch config can be written to a small JSON file with `store_patch_variables(...)` and read back with `load_patch_variables(...)`. The JSON captures `patch_size`, `target_spacing`, `apply_reorder`, the sampler settings, and the `normalization` spec, so an inference run can reconstruct the training config without re-typing it. This is the recommended way to ship a config alongside a checkpoint.\n", - "\n", - "The cell below writes the current config out and reads it back to confirm the round-trip. An inference run on another machine would simply start from `load_patch_variables(...)`." + "We load the five fold learners for `MODEL_KEY` into a list; each carries its own weights and preprocessing. `load_learner` uses Python `pickle`, which can execute arbitrary code, so only load checkpoints you trust." ] }, { "cell_type": "code", "execution_count": null, - "id": "39b66065", + "id": "4ce981f2", "metadata": {}, "outputs": [], "source": [ - "CONFIG_JSON = \"inference_patch_config.json\"\n", - "\n", - "# Coerce the normalization transforms to a JSON-serializable spec. PatchConfig does this\n", - "# coercion in __post_init__, so we build a throwaway config to obtain the spec list.\n", - "norm_specs = PatchConfig(patch_size=PATCH_SIZE, normalization=pre_inference_tfms).normalization\n", - "\n", - "# Persist the config so any inference run can reconstruct it verbatim.\n", - "store_patch_variables(\n", - " CONFIG_JSON,\n", - " patch_size=patch_config.patch_size,\n", - " patch_overlap=patch_config.patch_overlap,\n", - " aggregation_mode=patch_config.aggregation_mode,\n", - " apply_reorder=patch_config.apply_reorder,\n", - " target_spacing=patch_config.target_spacing,\n", - " sampler_type=patch_config.sampler_type,\n", - " label_probabilities=patch_config.label_probabilities,\n", - " samples_per_volume=patch_config.samples_per_volume,\n", - " queue_length=patch_config.queue_length,\n", - " queue_num_workers=patch_config.queue_num_workers,\n", - " keep_largest_component=patch_config.keep_largest_component,\n", - " normalization=norm_specs,\n", - ")\n", + "def _load_learner(pkl_path):\n", + " \"\"\"Load an exported fastai learner onto GPU if available, else CPU; set eval mode.\"\"\"\n", + " learn = load_learner(pkl_path, cpu=not torch.cuda.is_available())\n", + " learn.model.eval()\n", + " return learn\n", "\n", - "# Read it back to confirm the round-trip (inference elsewhere would start here).\n", - "reloaded = load_patch_variables(CONFIG_JSON)\n", - "print(\"Persisted config to\", CONFIG_JSON)\n", - "print(\"target_spacing:\", reloaded[\"target_spacing\"])\n", - "print(\"patch_size: \", reloaded[\"patch_size\"])\n", - "print(\"normalization: \", reloaded[\"normalization\"])" + "\n", + "# Load the fold learners; load_learner uses pickle, so only load checkpoints you trust.\n", + "assert FOLD_LEARNER_PATHS, (\n", + " \"No fold checkpoints found. Train the folds with notebook 01, or set \"\n", + " \"FOLD_LEARNER_PATHS = {1: '.../best_learner.pkl', ...} in the config cell.\")\n", + "\n", + "predictors = []\n", + "for fold in sorted(FOLD_LEARNER_PATHS):\n", + " predictors.append(_load_learner(FOLD_LEARNER_PATHS[fold]))\n", + " print(f\"Loaded fold {fold}: {FOLD_LEARNER_PATHS[fold]}\")\n", + "print(f\"Ensemble ready: {len(predictors)} fold model(s) for '{MODEL_KEY}'.\")" ] }, - { - "cell_type": "markdown", - "id": "20f7b65d", - "metadata": {}, - "source": "## 4. Load the models\n\nWe load the five fold learners for `MODEL_KEY` into a list. Each exported learner carries its own trained weights and preprocessing, so no architecture code is needed. `load_learner` uses Python `pickle`, which can execute arbitrary code, so only load checkpoints you trust.\n\nThe list of learners is what the inference engine soft-votes in the next section." - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4ce981f2", - "metadata": {}, - "outputs": [], - "source": "def _load_learner(pkl_path):\n \"\"\"Load an exported fastai learner onto GPU if available, else CPU; set eval mode.\"\"\"\n learn = load_learner(pkl_path, cpu=not torch.cuda.is_available())\n learn.model.eval()\n return learn\n\n\n# Load the five fold learners for MODEL_KEY into a list. Each exported learner carries its\n# own trained weights and preprocessing, so no architecture code is needed here. load_learner\n# uses Python pickle, which can execute arbitrary code, so only load checkpoints you trust.\nassert FOLD_LEARNER_PATHS, (\n \"No fold checkpoints found. Train the folds with notebook 01, or set \"\n \"FOLD_LEARNER_PATHS = {1: '.../best_learner.pkl', ...} in the config cell.\")\n\npredictors = []\nfor fold in sorted(FOLD_LEARNER_PATHS):\n predictors.append(_load_learner(FOLD_LEARNER_PATHS[fold]))\n print(f\"Loaded fold {fold}: {FOLD_LEARNER_PATHS[fold]}\")\nprint(f\"Ensemble ready: {len(predictors)} fold model(s) for '{MODEL_KEY}'.\")" - }, - { - "cell_type": "markdown", - "id": "ceaa05d1", - "metadata": {}, - "source": "### A note on SegMamba\n\nSegMamba is not part of MONAI; it comes from our fork **skaliy/SegMamba-V2**. Install it into this environment before selecting `MODEL_KEY = \"segmamba\"`:\n\n```bash\npip install -e /home/sathiesh/ml_projects/SegMamba-V2 # base (import path: models_segmamba)\npip install -e \"/home/sathiesh/ml_projects/SegMamba-V2[gpu]\" # adds mamba-ssm + causal-conv1d for CUDA\n```\n\nThe model is imported as `from models_segmamba.segmambav2 import SegMamba`. It supports two Mamba backends:\n\n- **`mamba_backend=\"mamba_ssm\"`** (default): the CUDA kernel used for training and fast GPU inference.\n- **`mamba_backend=\"mambamixer\"`**: a pure-PyTorch backend (needs `transformers`) that runs on CPU. It is weight-compatible with a `mamba_ssm`-trained checkpoint, so the same `.pth` loads with `strict=True` and no retraining.\n\nUNet and DynUNet need no fork and run on CPU or GPU as-is. One environment caveat: `causal-conv1d` and `mamba-ssm` must be compiled against the active PyTorch build, otherwise the import fails with an \"undefined symbol\" ABI error. On CPU we sidestep this with a small `_force_cpu_mamba_backend()` helper (defined in the section 8 cell below, mirroring `infer_segmamba_cpu.py`) that tells `transformers` not to probe the CUDA kernels. Section 8 shows the CPU path end to end." - }, { "cell_type": "markdown", "id": "37ea9e87", "metadata": {}, - "source": "## 5. Run inference\n\nEnsembling is built into fastMONAI: pass a **list of learners** to `patch_inference` and it\nsoft-votes them, averaging each patch's class probabilities across the models before taking the\nargmax, all in a single sliding-window pass. We pass the list of fold learners as\n`learner = predictors`.\n\nFor each scan `patch_inference` reorders, resamples, and normalizes the raw input exactly as in\ntraining, slides the patch grid, aggregates with the Hann window (soft-voting per patch across the\nfolds), resamples the prediction back to the input's original grid, and applies the training-time\npost-processing (argmax and keep-largest-component). Predictions are written to `OUTPUT_DIR` as\n`_pred.nii.gz` and returned as in-memory masks in `predictions` for the sections below.\n\nTwo options:\n\n- **`tta=USE_TTA`**: each patch is run through 8 axis-flip combinations and the probabilities\n averaged (combined with the fold averaging). Tighter borders, roughly 8x compute. Impractical on CPU.\n- **`amp=USE_AMP`**: float16 forward pass on CUDA. Faster and lighter; ignored on CPU.\n\nTo obtain the averaged probability map instead of a mask, pass `return_probabilities=True`." + "source": [ + "## 5. Run inference\n", + "\n", + "Pass a **list of learners** to `patch_inference` to soft-vote them: it averages each patch's class probabilities across the folds before the argmax, in one sliding-window pass. For each scan it reorders, resamples, and normalizes the input as in training, slides the Hann-blended patch grid, resamples the prediction back to the original grid, and post-processes (argmax + keep-largest-component). Predictions are written to `OUTPUT_DIR` as `_pred.nii.gz` and returned in `predictions`.\n", + "\n", + "- **`tta=USE_TTA`**: 8 axis-flip combinations averaged per patch. Tighter borders, ~8x compute, impractical on CPU.\n", + "- **`amp=USE_AMP`**: float16 forward pass; CUDA-only.\n", + "\n", + "Pass `return_probabilities=True` for the averaged probability map instead of a mask." + ] }, { "cell_type": "code", @@ -182,18 +199,44 @@ "id": "65abc50b", "metadata": {}, "outputs": [], - "source": "# Ensembling is a fastMONAI feature: a list of learners makes patch_inference soft-vote them\n# (per-patch probabilities averaged before argmax) in one sliding-window pass. Post-processing\n# (argmax, keep-largest-component) is applied to the averaged map.\nlearners = predictors\n\nPath(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)\npredictions = patch_inference(\n learner=learners,\n config=patch_config,\n file_paths=NEW_CASES,\n pre_inference_tfms=pre_inference_tfms,\n save_dir=OUTPUT_DIR,\n progress=True,\n tta=USE_TTA,\n amp=USE_AMP,\n)\n\n\ndef _pred_filename(input_path):\n \"\"\"Name patch_inference writes for a prediction: '_pred.nii[.gz]'. Mirrors the\n engine's saver, stripping only the trailing extension so '.nii' and '.nii.gz' inputs\n both map to the right output name (used again when we reload the mask to visualize).\"\"\"\n p = Path(input_path)\n if p.suffix == \".gz\" and p.stem.endswith(\".nii\"):\n return f\"{p.stem[:-4]}_pred.nii.gz\"\n if p.suffix == \".nii\":\n return f\"{p.stem}_pred.nii\"\n return f\"{p.stem}_pred.nii.gz\"\n\n\nprint(f\"\\nWrote {len(predictions)} prediction(s) to {OUTPUT_DIR}/ \"\n f\"(soft-vote ensemble of {len(predictors)} folds).\")\nfor p in NEW_CASES:\n print(f\" {p} -> {OUTPUT_DIR}/{_pred_filename(p)}\")" + "source": [ + "# A list of learners makes patch_inference soft-vote the folds; normalization comes from patch_config.\n", + "learners = predictors\n", + "\n", + "Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)\n", + "predictions = patch_inference(\n", + " learner=learners,\n", + " config=patch_config,\n", + " file_paths=NEW_CASES,\n", + " save_dir=OUTPUT_DIR,\n", + " progress=True,\n", + " tta=USE_TTA,\n", + " amp=USE_AMP,\n", + ")\n", + "\n", + "\n", + "def _pred_filename(input_path):\n", + " \"\"\"Output name patch_inference writes for a prediction: '_pred.nii[.gz]'.\"\"\"\n", + " p = Path(input_path)\n", + " if p.suffix == \".gz\" and p.stem.endswith(\".nii\"):\n", + " return f\"{p.stem[:-4]}_pred.nii.gz\"\n", + " if p.suffix == \".nii\":\n", + " return f\"{p.stem}_pred.nii\"\n", + " return f\"{p.stem}_pred.nii.gz\"\n", + "\n", + "\n", + "print(f\"\\nWrote {len(predictions)} prediction(s) to {OUTPUT_DIR}/ \"\n", + " f\"(soft-vote ensemble of {len(predictors)} folds).\")\n", + "for p in NEW_CASES:\n", + " print(f\" {p} -> {OUTPUT_DIR}/{_pred_filename(p)}\")" + ] }, { "cell_type": "markdown", "id": "691d43d1", "metadata": {}, "source": [ - "## 6. Visualize a prediction\n", - "\n", - "Now we inspect a result qualitatively. `vision_plot` is imported separately (it is not part of `vision_all`). We load the input scan as a `MedImage` and its predicted mask as a `MedMask`.\n", - "\n", - "Both are read with fastMONAI's default loader, which does not reorder or resample, so they share the input's original voxel grid and line up slice for slice. We pick the axial slice where the predicted tumor is largest with `find_max_slice`, then show the input, the mask, and an overlay. `voxel_size=TARGET_SPACING` sets only the display aspect ratio; it does not alter the data." + "## 6. Visualize a prediction" ] }, { @@ -202,7 +245,37 @@ "id": "ff90e7e6", "metadata": {}, "outputs": [], - "source": "from fastMONAI.vision_plot import * # fastMONAI plotting; imported separately from vision_all\nfrom torchio.visualization import rotate\nimport matplotlib.pyplot as plt\n\nidx = 0\nimg_fn = NEW_CASES[idx]\n# Same name patch_inference used when saving (see _pred_filename in the run-inference cell),\n# so the mask reloads correctly for both .nii and .nii.gz inputs.\npred_fn = Path(OUTPUT_DIR) / _pred_filename(img_fn)\n\n# Default fastMONAI loader: no reorder, no resample, so input and mask share the\n# input's original voxel grid and align slice for slice.\nimg = MedImage.create(img_fn)\npred_mask = MedMask.create(pred_fn)\n\nplane = 2 # 0 = sagittal, 1 = coronal, 2 = axial\nsl = int(find_max_slice(pred_mask.data[0].cpu().numpy(), plane))\nprint(f\"Predicted foreground voxels: {int(pred_mask.data.sum())}\")\nprint(f\"Showing plane={plane} (axial), slice={sl}\")\n\n\ndef _disp_slice(vol3d, i, plane, vsize):\n \"\"\"Replicate fastMONAI's show slicing so the overlay lines up with .show().\"\"\"\n sr, sa, ss = vsize\n ops = {0: (vol3d[i, :, :], ss / sa),\n 1: (vol3d[:, i, :], ss / sr),\n 2: (vol3d[:, :, i], sa / sr)}\n sl2d, aspect = ops[plane]\n return rotate(sl2d, radiological=True, n=1), aspect\n\n\nfig, axes = plt.subplots(1, 3, figsize=(15, 5))\nimg.show(ctx=axes[0], anatomical_plane=plane, slice_index=sl, voxel_size=TARGET_SPACING)\naxes[0].set_title(\"Input T1\")\npred_mask.show(ctx=axes[1], anatomical_plane=plane, slice_index=sl, voxel_size=TARGET_SPACING)\naxes[1].set_title(\"Predicted mask\")\n\nimg_slice, aspect = _disp_slice(img.data[0].cpu().numpy(), sl, plane, TARGET_SPACING)\nmsk_slice, _ = _disp_slice(pred_mask.data[0].cpu().numpy(), sl, plane, TARGET_SPACING)\naxes[2].imshow(img_slice, cmap=\"gray\", aspect=aspect)\naxes[2].imshow(np.ma.masked_where(msk_slice == 0, msk_slice),\n cmap=\"autumn\", alpha=0.5, aspect=aspect)\naxes[2].set_title(\"Overlay\")\naxes[2].axis(\"off\")\nplt.tight_layout()\nplt.show()" + "source": [ + "from fastMONAI.vision_plot import *\n", + "import matplotlib.pyplot as plt\n", + "\n", + "idx = 0\n", + "img_fn = NEW_CASES[idx]\n", + "# Same name patch_inference saved under, for both .nii and .nii.gz inputs.\n", + "pred_fn = Path(OUTPUT_DIR) / _pred_filename(img_fn)\n", + "\n", + "# Default loader: no reorder/resample, so input and mask share the original grid.\n", + "img = MedImage.create(img_fn)\n", + "pred_mask = MedMask.create(pred_fn)\n", + "\n", + "# Display aspect ratio only; does not alter data.\n", + "disp_spacing = patch_config.target_spacing\n", + "\n", + "plane = 2 # 0 = sagittal, 1 = coronal, 2 = axial\n", + "sl = int(find_max_slice(pred_mask.data[0].cpu().numpy(), plane))\n", + "print(f\"Predicted foreground voxels: {int(pred_mask.data.sum())}\")\n", + "print(f\"Showing plane={plane} (axial), slice={sl}\")\n", + "\n", + "fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n", + "img.show(ctx=axes[0], anatomical_plane=plane, slice_index=sl, voxel_size=disp_spacing)\n", + "axes[0].set_title(\"Input T1\")\n", + "pred_mask.show(ctx=axes[1], anatomical_plane=plane, slice_index=sl, voxel_size=disp_spacing)\n", + "axes[1].set_title(\"Predicted mask\")\n", + "show_mask_overlay(img, pred_mask, ctx=axes[2], anatomical_plane=plane,\n", + " slice_index=sl, voxel_size=disp_spacing, title=\"Overlay\")\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "markdown", @@ -211,15 +284,7 @@ "source": [ "## 7. Optional: validate against ground truth\n", "\n", - "If a case has an expert segmentation, we can score the prediction with the **same metric functions** notebook 01 uses in its cross-validation, so the numbers here are directly comparable. This cell is optional: it only runs when you set `GT_PATH` to a mask file.\n", - "\n", - "The metrics:\n", - "\n", - "- **DSC** (Dice) and **sensitivity** / **precision**: overlap-based agreement.\n", - "- **LDR** (lesion detection rate) and **signed RVE** (relative volume error).\n", - "- **Surface metrics in millimeters** via `calculate_surface_metrics`: ASSD, HD95, and NSD. These are spacing-aware, so we read the per-case voxel spacing straight from the ground-truth file with `tio.LabelMap(GT_PATH).spacing`. This matters because the VS cohort has non-uniform spacing, and using one global spacing would bias the surface distances.\n", - "\n", - "The prediction and ground truth are compared as 5D tensors `[batch, channel, X, Y, Z]`, matching the metric API." + "If a case has an expert segmentation, we score the prediction with the same metric functions notebook 01 uses, so the numbers are comparable. Runs only when `GT_PATH` points to a mask file." ] }, { @@ -234,8 +299,7 @@ "GT_PATH = None\n", "\n", "if GT_PATH:\n", - " # Score the first prediction with the library's per-case panel (loads the GT mask's\n", - " # data and spacing from one object).\n", + " # Score the first prediction with the library's per-case panel.\n", " row = evaluate_segmentations([predictions[0]], [GT_PATH]).iloc[0]\n", " print(f\"DSC: {row['dsc']:.4f}\")\n", " print(f\"Sensitivity: {row['sensitivity']:.4f}\")\n", @@ -255,17 +319,9 @@ "id": "df20f031", "metadata": {}, "source": [ - "## 8. Running SegMamba on CPU\n", - "\n", - "When no CUDA GPU is available, SegMamba can still run through the exact same `patch_inference` pipeline by switching to the `mambamixer` backend. This mirrors `research/vs_seg/infer_segmamba_cpu.py` and `research/vs_seg/SEGMAMBA_CPU.md`.\n", + "## 8. Running SegMamba on CPU (single-checkpoint demo)\n", "\n", - "The key points:\n", - "\n", - "- Construct with `mamba_backend=\"mambamixer\"` (pure PyTorch, needs `transformers`). No CUDA and no `mamba_ssm` required.\n", - "- A checkpoint trained with the default `mamba_ssm` backend loads into the `mambamixer` model unchanged (`strict=True`), so no retraining is needed. Correctness was verified decision-identical to GPU on a real case (Dice 1.0).\n", - "- It is slow: roughly 11 minutes per volume on CPU, and TTA is impractical there, so keep `tta=False`.\n", - "\n", - "The cell below is a self-contained CPU inference path. It is guarded so it does nothing unless `RUN_SEGMAMBA_CPU` is set to `True` and a checkpoint is present." + "This is a standalone demonstration of the CPU backend swap, not the ensemble path above. Only SegMamba needs it. Its Mamba blocks default to the `mamba_ssm`/`causal-conv1d` CUDA kernels, so on a CPU-only box you switch to `mamba_backend=\"mambamixer\"` (pure PyTorch, needs `transformers`). A checkpoint trained with the default `mamba_ssm` backend loads into it with `strict=True`, no retraining." ] }, { @@ -274,13 +330,52 @@ "id": "1eff224c", "metadata": {}, "outputs": [], - "source": "def _force_cpu_mamba_backend():\n \"\"\"Make transformers' MambaMixer take its pure-PyTorch path instead of probing CUDA\n kernels. A no-op on a clean CPU-only box; needed only when the environment has a\n broken or version-mismatched causal_conv1d / mamba_ssm build (the \"undefined symbol\"\n ABI error). Mirrors research/vs_seg/infer_segmamba_cpu.py.\"\"\"\n import sys\n import transformers # noqa: F401\n for modname in (\"transformers.utils.import_utils\", \"transformers.utils\"):\n mod = sys.modules.get(modname)\n if mod is not None:\n for fn in (\"is_causal_conv1d_available\", \"is_mamba_ssm_available\"):\n if hasattr(mod, fn):\n setattr(mod, fn, (lambda *a, **k: False))\n\n\nRUN_SEGMAMBA_CPU = False\nSEGMAMBA_WEIGHTS = \"models/best_segmamba.pth\"\n\nif RUN_SEGMAMBA_CPU:\n _force_cpu_mamba_backend() # skips the mismatched CUDA kernel probe (defined just above)\n from models_segmamba.segmambav2 import SegMamba\n\n # Pure-PyTorch Mamba backend: runs on CPU, no mamba_ssm / causal_conv1d needed.\n seg_model = SegMamba(\n in_chans=1, out_chans=2, depths=[2, 2, 2, 2],\n feat_size=[48, 96, 192, 384], hidden_size=768,\n mamba_backend=\"mambamixer\",\n )\n # A mamba_ssm-trained checkpoint loads strict into the mambamixer model unchanged.\n from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present\n sd = torch.load(SEGMAMBA_WEIGHTS, map_location=\"cpu\")\n sd = sd[\"model\"] if isinstance(sd, dict) and \"model\" in sd else sd\n consume_prefix_in_state_dict_if_present(sd, \"_orig_mod.\") # official torch helper\n seg_model.load_state_dict(sd, strict=True)\n seg_model.eval() # stays on CPU\n\n cpu_preds = patch_inference(\n learner=seg_model,\n config=patch_config,\n file_paths=NEW_CASES,\n pre_inference_tfms=pre_inference_tfms,\n save_dir=OUTPUT_DIR,\n progress=True,\n tta=False, # TTA is impractical on CPU\n )\n print(f\"SegMamba CPU: {len(cpu_preds)} prediction(s) written to {OUTPUT_DIR}/\")\nelse:\n print(\"Set RUN_SEGMAMBA_CPU = True (with a SegMamba checkpoint) to run this path.\")" - }, - { - "cell_type": "markdown", - "id": "53b15733", - "metadata": {}, - "source": "## Recap and pointers\n\n- **This notebook runs a 5-fold soft-voting ensemble.** Each fold's probability map is produced in the input's original space, the five are averaged, and only then decoded (argmax + keep-largest-component). This usually beats any single fold.\n- **Preprocessing parity is the rule that governs correctness.** The prediction is valid only because inference reused the training `apply_reorder`, `target_spacing`, and foreground `ZNormalization`. When in doubt, load an exported learner (which carries its own preprocessing) or ship a `PatchConfig` JSON via `store_patch_variables` / `load_patch_variables`.\n- **Patch inference** slides an overlapping grid, predicts per patch, and blends with a Hann window, so arbitrarily large volumes fit in memory.\n- **TTA** trades roughly 8x compute for tighter borders; **AMP** speeds up the GPU forward pass. Neither changes the required preprocessing.\n- **Model choice is orthogonal to this notebook.** UNet and DynUNet run anywhere; SegMamba needs the fork and prefers a GPU, with a documented CPU fallback.\n\nPointers:\n\n- Training and the checkpoints this notebook consumes: `01_five_fold_cross_validation.ipynb`.\n- An earlier whole-volume ensemble reference: `research/vs_seg/10_ensemble_inference_verification_unet.ipynb`.\n- SegMamba CPU details: `research/vs_seg/infer_segmamba_cpu.py` and `research/vs_seg/SEGMAMBA_CPU.md`.\n- The inference engine and config: `fastMONAI/vision_patch.py` (`patch_inference`, `PatchInferenceEngine`, `PatchConfig`)." + "source": [ + "def _force_cpu_mamba_backend():\n", + " \"\"\"Make transformers' MambaMixer take its pure-PyTorch path instead of probing CUDA kernels (which crash on CPU or on an ABI-mismatched mamba_ssm build).\"\"\"\n", + " import sys\n", + " import transformers # noqa: F401\n", + " for modname in (\"transformers.utils.import_utils\", \"transformers.utils\"):\n", + " mod = sys.modules.get(modname)\n", + " if mod is not None:\n", + " for fn in (\"is_causal_conv1d_available\", \"is_mamba_ssm_available\"):\n", + " if hasattr(mod, fn):\n", + " setattr(mod, fn, (lambda *a, **k: False))\n", + "\n", + "\n", + "RUN_SEGMAMBA_CPU = False\n", + "SEGMAMBA_WEIGHTS = \"models/best_segmamba.pth\" # single checkpoint on purpose: demos the CPU swap, not the fold ensemble\n", + "\n", + "if RUN_SEGMAMBA_CPU:\n", + " _force_cpu_mamba_backend() # skips the mismatched CUDA kernel probe\n", + " from models_segmamba.segmambav2 import SegMamba\n", + "\n", + " # Pure-PyTorch backend, runs on CPU.\n", + " seg_model = SegMamba(\n", + " in_chans=1, out_chans=2, depths=[2, 2, 2, 2],\n", + " feat_size=[48, 96, 192, 384], hidden_size=768,\n", + " mamba_backend=\"mambamixer\",\n", + " )\n", + " # A mamba_ssm-trained checkpoint loads strict into the mambamixer model unchanged.\n", + " from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present\n", + " sd = torch.load(SEGMAMBA_WEIGHTS, map_location=\"cpu\")\n", + " sd = sd[\"model\"] if isinstance(sd, dict) and \"model\" in sd else sd\n", + " consume_prefix_in_state_dict_if_present(sd, \"_orig_mod.\") # official torch helper\n", + " seg_model.load_state_dict(sd, strict=True)\n", + " seg_model.eval() # stays on CPU\n", + "\n", + " cpu_preds = patch_inference(\n", + " learner=seg_model,\n", + " config=patch_config,\n", + " file_paths=NEW_CASES,\n", + " save_dir=OUTPUT_DIR,\n", + " progress=True,\n", + " tta=False, # TTA is impractical on CPU\n", + " )\n", + " print(f\"SegMamba CPU: {len(cpu_preds)} prediction(s) written to {OUTPUT_DIR}/\")\n", + "else:\n", + " print(\"Set RUN_SEGMAMBA_CPU = True (with a SegMamba checkpoint) to run this path.\")" + ] } ], "metadata": { @@ -296,4 +391,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} diff --git a/research/vestibular_schwannoma/inference_patch_config.json b/research/vestibular_schwannoma/inference_patch_config.json deleted file mode 100644 index 366f441..0000000 --- a/research/vestibular_schwannoma/inference_patch_config.json +++ /dev/null @@ -1 +0,0 @@ -{"patch_size": [192, 192, 48], "patch_overlap": 0.5, "aggregation_mode": "hann", "apply_reorder": true, "target_spacing": [0.4102, 0.4102, 1.5], "sampler_type": "label", "label_probabilities": {"0": 0.2, "1": 0.8}, "samples_per_volume": 4, "queue_length": 1200, "queue_num_workers": 16, "keep_largest_component": true, "normalization": [{"name": "ZNormalization", "masking_method": "foreground", "channel_wise": true}]} \ No newline at end of file diff --git a/research/vestibular_schwannoma/research_pacs_integration/.dockerignore b/research/vestibular_schwannoma/research_pacs_integration/.dockerignore new file mode 100644 index 0000000..0aafdd1 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/.dockerignore @@ -0,0 +1,15 @@ +# Keep large/irrelevant files out of the Docker build context and image. +# NOTE: do NOT ignore the model directories or *.pkl weights - they must ship. + +# Saved docker image export(s) - e.g. vs-seg_docker_image_Sathiesh.tar.gz (~11.8 GB) +*.tar.gz +*.tar + +# Stale editor backup of the Dockerfile +.ror/virt/Dockerfile~ + +# Local tooling / caches +.claude/ +**/__pycache__/ +*.pyc +*.pyo diff --git a/research/vestibular_schwannoma/research_pacs_integration/.gitignore b/research/vestibular_schwannoma/research_pacs_integration/.gitignore new file mode 100644 index 0000000..4295205 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/.gitignore @@ -0,0 +1,13 @@ +# Fold weights + config + run-ids are generated on the build host by +# prepare_fold_models.py and are not committed (large binaries, and this folder +# is tracked). Regenerate before `docker build` - see README section 3. +vs5f_unet_models/ + +# Exported docker images (README section 4 writes these here) +*.tar.gz +*.tar + +# Python caches +__pycache__/ +*.pyc +*.pyo diff --git a/research/vestibular_schwannoma/research_pacs_integration/.ror/virt/Dockerfile b/research/vestibular_schwannoma/research_pacs_integration/.ror/virt/Dockerfile new file mode 100644 index 0000000..5b80838 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/.ror/virt/Dockerfile @@ -0,0 +1,47 @@ +FROM haukebartsch/fiona-component-python:latest +# +# build from ror-enabled directory as (conda_env must match the name in requirements.yml): +# docker build --no-cache --build-arg conda_env="fastmonai" -f .ror/virt/Dockerfile -t vs-seg-5fold:latest . +# + +WORKDIR /app + +# Set value to the name of your conda environment in requirements.yml (currently "fastmonai"). +# Optional: provide it to docker build as a variable +# --build-arg conda_env="fastmonai" +ARG conda_env=fastmonai + +LABEL "com.ror.vendor"="MMIV.no" \ + version="1.0" \ + com.ror.conda.env.name="${conda_env}" \ + description="Example docker container for workflow AI on the research information system at MMIV.no." + +ENV REPORT_FONT_PATH=/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf +ENV DCMDICTPATH=/dcmtk/dcmdata/data/dicom.dic:/dcmtk/dcmdata/data/private.dic + +# In case we need a proxy to connect to the internet on the build machine +#ENV HTTP_PROXY="http://proxy.ihelse.net:3128" +#ENV HTTPS_PROXY="http://proxy.ihelse.net:3128" + +# Create the environment: +COPY requirements.yml . +RUN conda env create -f requirements.yml + +# Make RUN commands use the new environment: +RUN echo "conda activate ${conda_env}" >> ~/.bashrc +SHELL ["/bin/bash", "--login", "-c"] + +# Demonstrate the environment is activated: +#RUN echo "Make sure fastMONAI is installed:" +#RUN python -c "import fastMONAI" + +# VERSION is a cheap runtime stamp (entrypoint.sh uses it for DICOM provenance). +# Declared after the conda layer so bumping it never invalidates that cache. +ARG VERSION=jul032026 +ENV VERSION=$VERSION + +COPY . /app + +# The code to run when container is started: +RUN ["chmod", "+x", "/app/entrypoint.sh"] +ENTRYPOINT [ "/app/entrypoint.sh" ] diff --git a/research/vestibular_schwannoma/research_pacs_integration/README.md b/research/vestibular_schwannoma/research_pacs_integration/README.md new file mode 100644 index 0000000..720dc0d --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/README.md @@ -0,0 +1,120 @@ +# UNet 5-Fold Ensemble ROR Container (Vestibular Schwannoma) + +## 1. What this is + +A ROR/PACS inference container for vestibular schwannoma segmentation that runs the five cross-validation UNet folds as a single soft-vote ensemble. `PatchInferenceEngine` receives the folds as a LIST and averages their per-patch softmax probabilities before the final argmax, so the folds behave as one logical model: the 5 folds averaged per patch. + +This container mirrors the production `vs_seg` container at `research/vs_seg/fastmonai_inference/integration`. Beyond model loading (a list of fold learners instead of one) and the provenance (ensemble run-ids), it also loads the full patch config from one file (section 3) rather than vs_seg's 3-key settings plus hardcoded values. Preprocessing, sliding-window inference and the output contract are otherwise the same. + +## 2. How it differs from vs_seg + +| Aspect | vs_seg (`research/vs_seg/fastmonai_inference/integration`) | This container | +|--------|-----------------------------------------------------------|----------------| +| Model loading | one exported learner (`unet_models/final_unet_learner.pkl`) | a LIST of fold learners (`vs5f_unet_models/fold_*.pkl`) | +| Averaging | none (single model) | per-patch softmax averaged by `PatchInferenceEngine` | +| SoftwareVersions | single-model tag plus one run-id | ensemble tag plus every fold run-id | +| Entrypoint INFO | `unet: ` | `unet N-fold ensemble` (N = shipped folds) | +| DICOM UID suffixes | `UNS1` / `UNP1` | `U5S1` / `U5P1` | +| display_name | `UNet` | `UNet 5-fold ensemble` | + +The distinct UID pair (`U5S1` segmentation, `U5P1` probability) means the ensemble series do not overwrite the vs_seg single-model series in PACS. + +Ensemble provenance is read from `vs5f_unet_models/mlflow_run_ids.txt` (one fold run-id per line, read with a guard for a missing file). `SoftwareVersions` becomes `[model_type + "-" + str(n) + "fold"]` (for this container, `unet-fold`) followed by the first 8 characters of each fold run-id and `"fastMONAI " + fastMONAI.__version__`, where `n` is the number of folds actually loaded. Apart from the documented changes (model loading, provenance, and the patch-config loading described in section 3), the container follows vs_seg. + +## 3. Prepare weights (before building) + +The fold weights are not committed to the repo (they are large, and unlike the untracked `vs_seg` tree this folder is tracked - see `.gitignore`). Generate `vs5f_unet_models/` on the build host before you build. In the `fastmonai-dev` conda env: + +```bash +source ~/miniconda3/etc/profile.d/conda.sh && conda activate fastmonai-dev +python prepare_fold_models.py +``` + +This writes three things into `vs5f_unet_models/`: + +- `fold_1.pkl .. fold_5.pkl` - the exported fold learners, copied AS-IS with no cleaning and no re-export, exactly like the vs_seg `unet_models/final_unet_learner.pkl`. +- `inference_patch_config.json` - the full patch-inference config, sourced from the canonical training config `../inference_patch_config.json`. +- `mlflow_run_ids.txt` - one fold run-id per line, used for DICOM provenance. + +Equivalent manual step (copy each exported fold `best_learner.pkl`, plus the config): + +```bash +mkdir -p vs5f_unet_models +cp /path/to/fold_1/best_learner.pkl vs5f_unet_models/fold_1.pkl +cp /path/to/fold_2/best_learner.pkl vs5f_unet_models/fold_2.pkl +# ... repeat for fold_3, fold_4, fold_5 +cp ../inference_patch_config.json vs5f_unet_models/inference_patch_config.json +``` + +`inference_patch_config.json` is read at inference with `load_patch_variables` and rebuilt into a full `PatchConfig(**config)`, so `patch_size`, `patch_overlap`, `aggregation_mode`, `apply_reorder`, `target_spacing`, `keep_largest_component` and `normalization` all come from the one training config - the single source of truth. A retrain that changes any of them then propagates to inference automatically (unlike vs_seg, which persists only `[patch_size, apply_reorder, target_spacing]` and hardcodes the rest in the stub). + +Note: only four UNet fold runs currently exist on disk, not a full five. The stub ensembles whatever N folds are present, so the container runs a 4-fold ensemble until the fifth fold is trained and prepared. + +## 4. Build, run, export + +These steps mirror `../../ROR_DOCKER_WORKFLOW.md` sections 3-7 with this container's values. Run every command from this integration folder. + +Pull the base image (unchanged from vs_seg): + +```bash +docker pull haukebartsch/fiona-component-python:latest +``` + +Build (Dockerfile `ARG VERSION=jul032026`): + +```bash +docker build --build-arg conda_env="fastmonai" -f .ror/virt/Dockerfile -t vs-seg-5fold:latest . +``` + +Run through ror trigger: + +```bash +ror trigger -cont vs-seg-5fold:latest -each -keep +``` + +Test-time augmentation (8-flip mirror) is on by default and improves accuracy, but it is much slower on CPU (~7-9x). Because this is an N-fold ensemble the cost stacks across folds: expect roughly 8-10 min/case without TTA and ~60-80 min/case with it for a 4-5 fold ensemble. Disable it with the `tta` key in `ROR_CONT_OPTIONS` (accepts true/false): + +```bash +ror trigger -cont vs-seg-5fold:latest -each -keep -envs '{"tta":false}' +``` + +Manual docker run (mount input read-only at `/data`, output at `/output`): + +```bash +docker run --rm \ + -v :/data:ro \ + -v :/output \ + vs-seg-5fold:latest +``` + +Export for distribution: + +```bash +docker save vs-seg-5fold:latest | gzip > vs-seg-5fold_docker_image_jul032026.tar.gz +``` + +See `../../ROR_DOCKER_WORKFLOW.md` for ror install, project init, and the full ror trigger option tables. + +## 5. Research PACS deployment + +Deployment mirrors vs_seg. The select statement is the same VS MR study selection; only the `docker_image` changes to `vs-seg-5fold:latest`. + +Select-statement entry: + +```json +"MMIVVestSchAI": { + "select": "SELECT study FROM study WHERE series named \"everything\" has Modality regexp \"MR\"", + "ROR_CONT_OPTIONS": "{}", + "docker_image": "vs-seg-5fold:latest" +} +``` + +Set `"ROR_CONT_OPTIONS": "{\"tta\":false}"` to disable TTA. The `config.json` trigger entry is unchanged from vs_seg. See `../../ROR_DOCKER_WORKFLOW.md` section 8 for the full fiona registration steps and the submission token. + +## 6. Output + +The container returns four series to PACS: `fused`, `fused_vote_map`, `reports`, and `mask`. The stub writes `/output/mask` (segmentation) and `/output/vote_map` (foreground probability); `entrypoint.sh` runs pr2mask, which produces `fused`, `fused_vote_map`, and `reports`. This is the same output contract as vs_seg. + +## 7. Caveat + +The exported `.pkl` fold learners are version-bound. The container's fastMONAI and fastai must match the versions used to train the folds, or the pickled learners will fail to load. diff --git a/research/vestibular_schwannoma/research_pacs_integration/entrypoint.sh b/research/vestibular_schwannoma/research_pacs_integration/entrypoint.sh new file mode 100755 index 0000000..b479c3f --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/entrypoint.sh @@ -0,0 +1,127 @@ +#!/bin/bash --login +# The --login ensures the bash configuration is loaded. + +output="/output" + +if [ ! -z "$CONDA_DEFAULT_ENV" ]; then + conda_env="$CONDA_DEFAULT_ENV" +fi + +if [ -z "$conda_env" ]; then + echo "Usage: " + exit -1 +fi + + +MODELTOUSE="unet" +USE_TTA=1 +# Handle environment variables provided to ror as --envs "ROR_CONT_OPTIONS={\"model-type\":\"unet\",\"tta\":false}" +if [ ! -z "$ROR_CONT_OPTIONS" ]; then + for key in $(echo "$ROR_CONT_OPTIONS" | jq -r "keys[]"); do + value=$(echo "$ROR_CONT_OPTIONS" | jq -r '."'$key'"') + if [ $key == "model-type" ]; then + MODELTOUSE=$value + fi + # 8-flip test-time augmentation is on by default; tta false/0/no/off + # disables it (true/1/yes/on re-enables). + if [ $key == "tta" ]; then + case "$(echo "$value" | tr '[:upper:]' '[:lower:]')" in + true|1|yes|on) USE_TTA=1 ;; + false|0|no|off) USE_TTA=0 ;; + esac + fi + done +fi + + +# This container ships a single unet 5-fold ensemble (model-type unet). TTA is on +# by default; set the tta key in ROR_CONT_OPTIONS to false to disable it. +SCRIPT_DIR="$(dirname "$0")" +STUB_SCRIPT="${SCRIPT_DIR}/stub_inference.py" +MODEL_ARG="--model-type ${MODELTOUSE}" +TTA_ARG="" +if [ "$USE_TTA" -eq 1 ]; then + TTA_ARG="--tta" + echo "Using ${MODELTOUSE} model (TTA: on)" +else + echo "Using ${MODELTOUSE} model (TTA: off)" +fi + +# where is pr2mask? +export PATH="/pr2mask:$PATH" + +# Fold count for the DICOM info field. Count the fold_*.pkl actually shipped in +# vs5f_unet_models/ (the same files the stub globs and loads), so this label and +# the stub's SoftwareVersions tag (len(models)) always agree. nullglob makes a +# missing directory / no matches yield 0 rather than a literal glob pattern. +shopt -s nullglob +fold_pkls=("${SCRIPT_DIR}"/vs5f_unet_models/fold_*.pkl) +shopt -u nullglob +n_folds=${#fold_pkls[@]} +INFO="${MODELTOUSE} ${n_folds}-fold ensemble, Predicted $(date '+%b%d%Y')" + +# if we find imageAndMask2Report and json2SR in this container +auto_report_mode=0 +output2="/output_tmp" +if [ -f /pr2mask/imageAndMask2Report ]; then + auto_report_mode=1 +else + output2="${output}" +fi + +# relax strict mode for conda activate (its hooks aren't 'set -euo pipefail' safe), then restore +set +euo pipefail +conda activate "${conda_env}" +if [ $? -ne 0 ]; then + echo "Error: activating conda environment \"$1\" failed." + exit -1 +fi +set -euo pipefail + +log_file="${output}"/stub_command.log +# prefer the bundled inference script; otherwise run the command passed to the container +if [ -n "$STUB_SCRIPT" ] && [ -f "$STUB_SCRIPT" ]; then + cmd="python ${STUB_SCRIPT} /data ${output2} ${MODEL_ARG} ${TTA_ARG}" +else + cmd="$@ ${output2}" +fi +echo "run now: $cmd" +# eval splits the assembled $cmd into command + arguments +eval $cmd + +if [ "$auto_report_mode" -eq 1 ]; then + echo "imageAndMask2Report:" + /pr2mask/imageAndMask2Report /data/input "${output2}/mask" "${output2}" -u "$VERSION" -i "$VERSION" --reporttype mosaic -t "${INFO} " >> "${log_file}" 2>&1 + echo "imageAndMask2Fused:" + /pr2mask/imageAndMask2Fused /data/input "${output2}/mask" "${output2}" -u "${VERSION}_fused" -i "$VERSION" >> "${log_file}" 2>&1 + echo "imageAndMask2Fused (vote map):" + /pr2mask/imageAndMask2Fused /data/input "${output2}/vote_map" "${output2}" --votemapmax 65535 --votemapagree 0.5 -u "${VERSION}_votemap" -s "peak agreement {peak_agreement}" -i "$VERSION" >> "${log_file}" 2>&1 + + # Four DICOM series are sent back to PACS: the report (reports/), the raw + # segmentation mask (mask/), the fused mask overlay (fused/) and the fused + # vote-map / agreement overlay (fused_vote_map/). The raw probability map + # (vote_map/), the B&W labels, the redcap JSON and the structured report + # (*.dcm) are intentionally not copied. + if [ ! -d "${output2}"/fused ]; then + echo "Error: no /fused folder found in ${output2}" + fi + cp -R "${output2}"/fused "${output}" + + if [ ! -d "${output2}"/fused_vote_map ]; then + echo "Error: no /fused_vote_map folder found in ${output2}" + fi + cp -R "${output2}"/fused_vote_map "${output}" + + if [ ! -d "${output2}"/reports ]; then + echo "Error: no /report folder found in ${output2}" + fi + cp -R "${output2}"/reports "${output}" + + if [ ! -d "${output2}"/mask ]; then + echo "Error: no /mask folder found in ${output2}" + fi + cp -R "${output2}"/mask "${output}" + + chmod -R 777 /output +fi +echo "$(date): processing done" >> "${log_file}" diff --git a/research/vestibular_schwannoma/research_pacs_integration/prepare_fold_models.py b/research/vestibular_schwannoma/research_pacs_integration/prepare_fold_models.py new file mode 100644 index 0000000..04af3f4 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/prepare_fold_models.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +''' +Bundle exported cross-validation fold learners into vs5f_unet_models/ so the +UNet 5-fold ensemble container can be built. + +Each exported best_learner.pkl is copied as-is: no loading, no re-export, and no +torch.compile unwrap. The raw exported learners load and run in the container, +matching the single-model vs_seg build that already runs in production. + +Run once in the fastmonai-dev env, for example: + python prepare_fold_models.py --experiment vs5f_unet +''' + +import argparse +import glob +import shutil +from pathlib import Path + +from fastMONAI.vision_all import * + + +def run_id_from_path(path): + """Parse the MLflow run-dir hash from an mlruns-style artifact path. + + Expects a path of the form ...//artifacts/... and returns , + or None when the path does not follow that layout. + """ + parts = Path(path).parts + if "artifacts" in parts: + idx = parts.index("artifacts") + if idx > 0: + return parts[idx - 1] + return None + + +def discover_learners(args): + """Return an ordered list of (src_path, run_id) for the fold learners. + + Explicit --pkl-paths win and are used in the given order. Otherwise try + find_fold_learners against a live MLflow store, then fall back to the on-disk + mlruns glob. The MLflow attempt is guarded so a missing store degrades to the + glob instead of crashing. + """ + if args.pkl_paths: + return [(p, run_id_from_path(p) or "unknown") for p in args.pkl_paths] + + try: + found = find_fold_learners(args.experiment) + except Exception as err: + print(f"WARNING: MLflow discovery failed ({err}); using on-disk mlruns glob.") + found = {} + + if found: + return [(src, run_id_from_path(src) or str(fold)) + for fold, src in sorted(found.items())] + + matches = sorted(glob.glob(args.mlruns_glob)) + if matches: + print("WARNING: no MLflow fold learners found; falling back to on-disk " + "mlruns glob. On-disk mlruns is tag-stripped, so fold order follows " + "path order and may not match the training fold numbers. Prefer " + "--pkl-paths (explicit fold order) or a live MLflow store.") + return [(src, run_id_from_path(src) or "unknown") for src in matches] + + +def main(): + parser = argparse.ArgumentParser( + prog="PrepareFoldModels", + description="Copy exported cross-validation fold learners into the " + "ensemble models directory (no cleaning or re-export)." + ) + parser.add_argument("--experiment", type=str, default="vs5f_unet", + help="MLflow experiment to discover fold learners from " + "(default: vs5f_unet).") + parser.add_argument("--out", type=Path, + default=Path(__file__).parent / "vs5f_unet_models", + help="Output models directory (default: ./vs5f_unet_models).") + parser.add_argument("--pkl-paths", nargs="*", default=None, + help="Explicit exported best_learner.pkl paths in fold " + "order; overrides discovery.") + parser.add_argument("--mlruns-glob", type=str, + default="mlruns/1/*/artifacts/model/best_learner.pkl", + help="Fallback glob for on-disk mlruns learners.") + parser.add_argument("--patch-size", nargs=3, type=int, default=[192, 192, 48], + help="Fallback patch size if --patch-config is missing " + "(default: 192 192 48).") + parser.add_argument("--target-spacing", nargs=3, type=float, + default=[0.4102, 0.4102, 1.5], + help="Fallback voxel spacing if --patch-config is missing " + "(default: 0.4102 0.4102 1.5).") + parser.add_argument("--reorder", default=True, action=argparse.BooleanOptionalAction, + help="Reorder to RAS+ canonical orientation (default: True).") + parser.add_argument("--patch-config", type=Path, + default=Path(__file__).parent.parent / "inference_patch_config.json", + help="Canonical training patch config to bundle for inference " + "(default: ../inference_patch_config.json). When present it is " + "the single source of truth; the --patch-size / --target-spacing " + "/ --reorder args are only the fallback used if it is missing.") + args = parser.parse_args() + + out = args.out + out.mkdir(parents=True, exist_ok=True) + + folds = discover_learners(args) + if not folds: + raise SystemExit( + "No fold learners found. Pass --pkl-paths, run against a live MLflow " + "store, or point --mlruns-glob at exported best_learner.pkl files.") + + # Copy each exported learner as-is (no loading, no re-export, no cleaning). + run_ids = [] + for i, (src, run_id) in enumerate(folds, start=1): + dst = out / ("fold_" + str(i) + ".pkl") + shutil.copy(src, dst) + run_ids.append(run_id) + print(f"Copied fold {i}: {src} -> {dst} (run {run_id})") + + # Ensemble provenance, one run id per line in fold_i.pkl order. + run_ids_path = out / "mlflow_run_ids.txt" + run_ids_path.write_text("\n".join(run_ids) + "\n") + + # Full patch-inference config, read back at inference with load_patch_variables + # into PatchConfig(**config). Prefer the canonical training config so every + # patch/preprocessing/postprocessing parameter (overlap, aggregation, + # normalization, keep_largest_component, ...) stays coupled to training; fall + # back to the CLI args + known-training defaults only if it is missing. + settings_path = out / "inference_patch_config.json" + if args.patch_config and args.patch_config.is_file(): + config = load_patch_variables(args.patch_config) + print(f"Sourced patch config from {args.patch_config}") + else: + print(f"WARNING: {args.patch_config} not found; falling back to CLI/default " + "patch config. Verify these match training.") + config = dict( + patch_size=args.patch_size, + patch_overlap=0.5, + aggregation_mode="hann", + apply_reorder=args.reorder, + target_spacing=args.target_spacing, + keep_largest_component=True, + normalization=[{"name": "ZNormalization", + "masking_method": "foreground", + "channel_wise": True}], + ) + store_patch_variables(str(settings_path), **config) + + print("\n" + "=" * 60) + print(f"Bundled {len(folds)} fold learners into {out}") + print(f"Run ids: {run_ids}") + print(f"Settings: {settings_path}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/research/vestibular_schwannoma/research_pacs_integration/requirements.yml b/research/vestibular_schwannoma/research_pacs_integration/requirements.yml new file mode 100644 index 0000000..3a5a0a7 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/requirements.yml @@ -0,0 +1,16 @@ +name: fastmonai +channels: + - pytorch + - conda-forge + - nodefaults +dependencies: + - python=3.11 + - pytorch=2.6.0 + - torchvision + - torchaudio + - cpuonly + - pip + - pip: + - fastMONAI==0.9.2 + - SimpleITK + - ipython \ No newline at end of file diff --git a/research/vestibular_schwannoma/research_pacs_integration/stub_inference.py b/research/vestibular_schwannoma/research_pacs_integration/stub_inference.py new file mode 100644 index 0000000..6219d67 --- /dev/null +++ b/research/vestibular_schwannoma/research_pacs_integration/stub_inference.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +''' +5-fold soft-vote ensemble patch-based inference for vestibular schwanoma segmentation. +./stub_inference.py --model-type [unet] + +Loads the five cross-validation fold learners (fold_*.pkl) and runs patch-based +sliding-window inference as one ensemble. PatchInferenceEngine averages the +per-patch class probabilities across the folds before argmax, so the folds +soft-vote into a single segmentation. Each fold is an exported fastai learner, +loaded as-is with load_learner. +''' + +import argparse +import glob +import json +from pathlib import Path + +import numpy as np +from fastai.learner import load_learner +import fastMONAI +from fastMONAI.vision_all import * +from fastMONAI.vision_patch import PatchConfig, PatchInferenceEngine +from fastMONAI.vision_inference import keep_largest +from imagedata.series import Series + + +# Model configuration (single model type; five folds soft-voted as one ensemble) +MODEL_CONFIGS = { + "unet": { + "models_dir": "vs5f_unet_models", + "weights_glob": "fold_*.pkl", + "seg_uid": "U5S1", + "prob_uid": "U5P1", + "display_name": "UNet 5-fold ensemble", + }, +} + +# Path setup - models are relative to script location +SCRIPT_DIR = Path(__file__).parent + +# Use CPU for inference +USE_CPU = True + +# Runtime batching knob (not a PatchConfig field). Every patch/preprocessing +# parameter comes from the persisted patch config; see run_inference. +SW_BATCH_SIZE = 2 + + +def save_series_pred(series_obj, save_dir, val='1234'): + """Save series prediction with updated DICOM UIDs. + + Makes sure we get derived UIDs to allow for overwrite of image objects in PACS. + """ + # Update series instance UID + my_seriesInstanceUID = series_obj.seriesInstanceUID[:-4] + val + series_obj.seriesInstanceUID = my_seriesInstanceUID + + # Update study ID from patient ID + if hasattr(series_obj, 'patientID') and series_obj.patientID: + my_studyID = series_obj.patientID[3:] if len(series_obj.patientID) > 3 else series_obj.patientID + series_obj.studyID = my_studyID + + base_sop_uid = series_obj.getDicomAttribute('SOPInstanceUID') + + for slice_idx in range(series_obj.slices): + # Create unique UID per slice by modifying the base UID + my_SOPInstanceUID = base_sop_uid[:-8] + val + str(slice_idx).zfill(4) + series_obj.setDicomAttribute('SOPInstanceUID', my_SOPInstanceUID, slice=slice_idx) + + series_obj.write(save_dir, opts={'keep_uid': True}, formats=['dicom']) + + +def create_dicom_mask(pred, dicom_input_path, output_dir, uid_suffix, software_versions): + """ + Create DICOM series from segmentation mask using template series metadata. + + `pred` is the prediction in original image space (a tensor/array, already + resized and reoriented by PatchInferenceEngine.predict()). + """ + # Create a copy of the template series to hold the mask + mask_obj = Series(str(dicom_input_path), opts={'slice_tolerance': 1e-2}) + + # Overwrite SoftwareVersions (0018,1020) with the ensemble provenance passed in + # (an ensemble tag, the per-fold run ids, and the fastMONAI version); any value + # carried over from the source series is replaced. + mask_obj.setDicomAttribute('SoftwareVersions', software_versions) + + # Mark the mask as a derived/secondary object in ImageType (0008,0008): + # value 1 -> DERIVED, value 2 -> SECONDARY, keep any trailing values, and + # append a MASK marker (e.g. ORIGINAL\PRIMARY\M\ND -> DERIVED\SECONDARY\M\ND\MASK). + image_type = mask_obj.getDicomAttribute('ImageType') + image_type = [] if image_type is None else ( + [image_type] if isinstance(image_type, str) else list(image_type)) + mask_obj.setDicomAttribute('ImageType', ['DERIVED', 'SECONDARY'] + image_type[2:] + ['MASK']) + + # Get mask data and transform to DICOM orientation + new_mask = pred.numpy() + new_mask = new_mask.squeeze() + + # Transform to DICOM orientation + new_mask = np.transpose(new_mask, (-1, 1, 0)) + new_mask = new_mask.copy() # Make contiguous + new_mask = new_mask.astype(np.uint16) + + # Assign mask data to series using slice assignment + mask_obj[:] = new_mask + + # Create output directory + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Save with proper UID handling + save_series_pred(mask_obj, str(output_path), val=uid_suffix) + + return output_path + + +def create_dicom_prob_mask(pred, dicom_input_path, output_dir, uid_suffix): + """ + Create DICOM series from probability mask. + Scales probabilities from [0,1] to uint16 [0,65535]. + + `pred` is the foreground probability in original image space (a tensor/array). + """ + # Create a copy of the template series to hold the probability mask + mask_obj = Series(str(dicom_input_path), opts={'slice_tolerance': 1e-2}) + + # Get probability data and transform to DICOM orientation + prob_data = pred.numpy().squeeze() + + # Transform to DICOM orientation + prob_data = np.transpose(prob_data, (-1, 1, 0)) + prob_data = prob_data.copy() # Make contiguous + + # Scale from [0,1] to [0,65535] for uint16 + prob_scaled = (prob_data * 65535).astype(np.uint16) + + # Assign scaled probability data to series + mask_obj[:] = prob_scaled + + # Create output directory + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Save with proper UID handling + save_series_pred(mask_obj, str(output_path), val=uid_suffix) + + return output_path + + +def load_models(config): + """Load the exported fold learners and return their eval-mode networks as a list. + + Each fold learner (fold_*.pkl) is an exported fastai learner, loaded as-is with + load_learner. The returned list is handed to PatchInferenceEngine, which + soft-votes the folds (per-patch probability averaging) at inference time. + """ + paths = sorted(glob.glob(str(SCRIPT_DIR / config["models_dir"] / config["weights_glob"]))) + assert paths, ( + f"No fold_*.pkl found in {config['models_dir']}; " + "copy the exported fold learners there " + "(see README / prepare_fold_models.py)") + + print(f"\nLoading {config['display_name']} folds from: " + f"{SCRIPT_DIR / config['models_dir']}") + models = [] + for path in paths: + learn = load_learner(path, cpu=True) + model = learn.model + model.eval() + models.append(model) + print(f" Loaded fold: {path}") + + if len(paths) != 5: + print(f" WARNING: expected 5 folds, loaded {len(paths)}") + + return models + + +def build_software_versions(models_dir, model_type, n_folds): + """Build the SoftwareVersions values recording the ensemble provenance. + + Reads one MLflow run id per line from mlflow_run_ids.txt (guarded so a + missing or unreadable file never aborts inference) and returns the model + tag, the per-fold run-id prefixes, and the fastMONAI version. + """ + try: + run_ids = [line.strip() for line + in (models_dir / "mlflow_run_ids.txt").read_text().splitlines() + if line.strip()] + except OSError: + run_ids = [] + return ([f"{model_type}-{n_folds}fold"] + + [rid[:8] for rid in run_ids] + + [f"fastMONAI {fastMONAI.__version__}"]) + + +def run_inference(datafolder, output, model_type, use_tta=False): + """Run 5-fold soft-vote ensemble patch-based inference for the model type. + + use_tta enables 8-flip mirror test-time augmentation (default off). + """ + config = MODEL_CONFIGS[model_type] + display_name = config["display_name"] + + print("=" * 60) + print(f"{display_name} Inference - Vestibular Schwanoma Segmentation") + print("=" * 60) + + models_dir = SCRIPT_DIR / config["models_dir"] + + # Load the fold models first so a missing or empty models directory fails with an + # actionable message (see load_models) before anything else runs. Passing the list + # of models makes PatchInferenceEngine soft-vote them (per-patch probability + # averaging before argmax). The engine also handles preprocessing (reorder/resample + # via med_img_reader + Z-normalization), sliding-window patch sampling, + # resize-to-original, and reorientation. + models = load_models(config) + + # Load the full patch-inference config saved next to the models. This is the + # single source of truth for preprocessing and patch settings (patch_size, + # patch_overlap, aggregation_mode, apply_reorder, target_spacing, + # keep_largest_component, normalization), so training and inference stay in + # sync. PatchConfig(**config) reconstructs it; the engine reads apply_reorder, + # target_spacing and normalization from the config. + settings_path = models_dir / "inference_patch_config.json" + print(f"\nLoading inference settings from: {settings_path}") + patch_config = PatchConfig(**load_patch_variables(settings_path)) + print(f" Patch size: {patch_config.patch_size}") + print(f" Reorder: {patch_config.apply_reorder}") + print(f" Resample: {patch_config.target_spacing}") + + # Ensemble provenance for the mask DICOM (SoftwareVersions tag). + software_versions = build_software_versions(models_dir, model_type, len(models)) + + engine = PatchInferenceEngine(models, patch_config, sw_batch_size=SW_BATCH_SIZE) + + # predict() reads the DICOM folder directly and returns class probabilities + # in the original image space (resized + reoriented). The result is + # [n_classes, *spatial], though the released reorientation path adds a leading + # singleton ([1, n_classes, *spatial]); squeeze(0) normalizes both (it only + # drops dim 0 when its size is 1, and n_classes here is 2). The binary mask is + # derived via argmax + largest-connected-component, matching the validated CV path. + print(f"\nRunning patch-based inference " + f"(patch={patch_config.patch_size}, overlap={patch_config.patch_overlap}, " + f"agg={patch_config.aggregation_mode}, " + f"TTA={'on' if use_tta else 'off'})...") + prob = engine.predict(datafolder, return_probabilities=True, tta=use_tta).squeeze(0) + tumor_prob = prob[1] + segmentation = keep_largest(prob.argmax(0).float()) + + # Save outputs + print(f"\nSaving outputs to: {output}") + + # Save binary segmentation mask + mask_output_dir = output + '/mask' + create_dicom_mask(segmentation, datafolder, mask_output_dir, + uid_suffix=config["seg_uid"], software_versions=software_versions) + print(f" Saved segmentation mask to: {mask_output_dir}") + + # Save probability mask + prob_output_dir = output + '/vote_map' + create_dicom_prob_mask(tumor_prob, datafolder, prob_output_dir, uid_suffix=config["prob_uid"]) + print(f" Saved probability mask to: {prob_output_dir}") + + # Optional description passthrough (guarded so a missing/foreign descr.json + # never aborts inference after the masks have been written) + description = {} + input_description_path = Path(datafolder + "/descr.json") + if input_description_path.is_file(): + try: + with open(input_description_path, "r") as file: + description = json.load(file) + if isinstance(description, list) and description and isinstance(description[0], dict): + description[0]["ProbabilityMask"] = prob_output_dir + except (json.JSONDecodeError, OSError) as err: + print(f" Warning: could not read descr.json ({err})") + + # Summary + seg_array = segmentation.cpu().numpy().squeeze() + tumor_voxels = np.sum(seg_array > 0) + total_voxels = seg_array.size + + print("\n" + "=" * 60) + print(f"{display_name.upper()} INFERENCE COMPLETE") + print("=" * 60) + print(f"Folds ensembled: {len(models)}") + print(f"TTA: {'on' if use_tta else 'off'}") + print(f"Tumor voxels: {int(tumor_voxels)}") + print(f"Total voxels: {total_voxels}") + print(f"Tumor percentage: {100 * tumor_voxels / total_voxels:.4f}%") + print("=" * 60) + + +def main(): + # Argument parsing + parser = argparse.ArgumentParser( + prog='VestibularSchwanomaSegmentation', + description='Vestibular schwanoma 5-fold ensemble segmentation script.' + ) + parser.add_argument('fn', type=str, help='Directory name of the input folder') + parser.add_argument('on', type=str, help='Directory name for the output') + parser.add_argument('--model-type', type=str, choices=['unet'], + default='unet', help='Model type to use (default: unet)') + parser.add_argument('--tta', action='store_true', + help='Enable 8-flip mirror test-time augmentation ' + '(default: off; much slower on CPU, ~7-9x).') + args = parser.parse_args() + + datafolder = args.fn + '/input' + output = args.on + + run_inference(datafolder, output, args.model_type, use_tta=args.tta) + + +if __name__ == "__main__": + main()