diff --git a/research/vestibular_schwannoma/README.md b/research/vestibular_schwannoma/README.md index 38c3c5c..c108e79 100644 --- a/research/vestibular_schwannoma/README.md +++ b/research/vestibular_schwannoma/README.md @@ -6,11 +6,12 @@ training, inference on new cases, and PACS deployment. ## Contents +- `train_5fold.py`: command-line five-fold training and evaluation. - `notebooks/01_five_fold_cross_validation.ipynb`: train and compare UNet, DynUNet, and optional SegMamba models. - `notebooks/02_inference_new_cases.ipynb`: run one declared model or an explicit ensemble. - `workflow/`: project-local configuration, training model definitions, result aggregation, - and inference artifact handling used by the notebooks. + and inference artifact handling shared by the CLI and notebooks. - `data/ml_dataset.csv`: public case index and fixed fold assignments. - `deployment/pacs/`: Safetensors bundle builder and ROR/PACS container. - `tests/workflow/`: CPU-only workflow contract and orchestration tests. @@ -32,7 +33,18 @@ optional SegMamba fork when needed. ## Run -Start Jupyter from this directory or `notebooks/`: +Use the CLI for unattended training. Examples for a quick check, one complete model, and the +full comparison: + +```bash +python train_5fold.py --models unet --folds 1 --epochs 5 --no-compile +python train_5fold.py --models unet # One model, all five folds +python train_5fold.py --skip-unavailable +``` + +The default requests four models across five folds for 500 epochs; run +`python train_5fold.py --help` before starting. For interactive inspection and visualizations, +start Jupyter from this directory or `notebooks/`: ```bash jupyter lab notebooks/01_five_fold_cross_validation.ipynb diff --git a/research/vestibular_schwannoma/tests/workflow/test_config.py b/research/vestibular_schwannoma/tests/workflow/test_config.py index d3de687..d07499e 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_config.py +++ b/research/vestibular_schwannoma/tests/workflow/test_config.py @@ -22,6 +22,7 @@ def test_defaults_preserve_the_notebook_experiment(self): self.assertEqual(config.training_seed, 42) self.assertEqual(config.queue_num_workers, 4) self.assertEqual(config.queue_length, 300) + self.assertEqual(config.foreground_sampling_probability, 0.8) self.assertTrue(config.use_tta) def test_configuration_is_frozen(self): @@ -44,6 +45,14 @@ def test_invalid_declarations_fail_early(self): {"model_keys": ("unet",), "training_seed": -1}, {"model_keys": ("unet",), "training_seed": True}, {"model_keys": ("unet",), "patch_size": (192, 192, 0)}, + { + "model_keys": ("unet",), + "foreground_sampling_probability": 0, + }, + { + "model_keys": ("unet",), + "foreground_sampling_probability": 1, + }, ] for kwargs in cases: with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): @@ -71,6 +80,14 @@ def test_patch_factory_preserves_training_and_inference_contract(self): }, ) + def test_patch_factory_supports_70_30_sampling(self): + config = ExperimentConfig( + model_keys=("unet",), foreground_sampling_probability=0.7 + ) + patch = make_patch_config(config, []) + self.assertAlmostEqual(patch.label_probabilities[0], 0.3) + self.assertAlmostEqual(patch.label_probabilities[1], 0.7) + if __name__ == "__main__": unittest.main() diff --git a/research/vestibular_schwannoma/tests/workflow/test_models.py b/research/vestibular_schwannoma/tests/workflow/test_models.py index 4076b9e..9a1d6e7 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_models.py +++ b/research/vestibular_schwannoma/tests/workflow/test_models.py @@ -8,10 +8,15 @@ class TrainingModelConfigTests(unittest.TestCase): def test_specs_are_the_single_architecture_declaration(self): self.assertEqual(models.UNET_SPEC["arch_id"], "monai.unet") self.assertEqual(models.DYNUNET_SPEC["arch_id"], "monai.dynunet") + self.assertEqual(models.DYNUNET_SMALL_SPEC["arch_id"], "monai.dynunet") self.assertEqual( models.DYNUNET_SPEC["wrapper_spec"][0]["wrapper_id"], "fastmonai.dynunet_ds_adapter", ) + self.assertEqual( + models.DYNUNET_SMALL_SPEC["arch_kwargs"]["filters"], + [32, 64, 128, 256, 512], + ) self.assertEqual(models.SEGMAMBA_SPEC["arch_id"], "segmamba.v2") def test_declared_order_is_preserved(self): diff --git a/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py b/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py new file mode 100644 index 0000000..7ed2edf --- /dev/null +++ b/research/vestibular_schwannoma/tests/workflow/test_train_5fold.py @@ -0,0 +1,27 @@ +import unittest + +from vestibular_schwannoma import train_5fold + + +class FiveFoldLauncherTests(unittest.TestCase): + def test_defaults_select_all_models_and_preserve_80_20_control(self): + args = train_5fold._parser().parse_args([]) + + self.assertEqual(tuple(args.models), train_5fold.DEFAULT_MODELS) + self.assertEqual(args.folds, [1, 2, 3, 4, 5]) + self.assertEqual(args.foreground_probability, 0.8) + self.assertNotIn("use_tta", vars(args)) + + def test_70_30_sampling_remains_an_explicit_comparison(self): + args = train_5fold._parser().parse_args(["--foreground-probability", "0.7"]) + + self.assertEqual(args.foreground_probability, 0.7) + + def test_small_dynunet_uses_the_intermediate_512_bottleneck(self): + spec = train_5fold.TRAINING_MODEL_CONFIGS["dynunet_small"].model_spec + + self.assertEqual(spec["arch_kwargs"]["filters"], [32, 64, 128, 256, 512]) + + +if __name__ == "__main__": + unittest.main() diff --git a/research/vestibular_schwannoma/tests/workflow/test_training.py b/research/vestibular_schwannoma/tests/workflow/test_training.py index b8193d6..c58212b 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_training.py +++ b/research/vestibular_schwannoma/tests/workflow/test_training.py @@ -232,6 +232,7 @@ def test_fold_tracking_uses_fixed_model_and_output_contract(self): ) extra_params = create_callback.call_args.kwargs["extra_params"] self.assertEqual(extra_params["training_seed"], 42) + self.assertEqual(extra_params["foreground_sampling_probability"], 0.8) self.assertEqual( json.loads(extra_params["gpu_augmentations"]), [ @@ -332,6 +333,7 @@ def build_model(*args, **kwargs): ) extra_params = create_callback.call_args.kwargs["extra_params"] self.assertEqual(extra_params["training_seed"], 42) + self.assertEqual(extra_params["foreground_sampling_probability"], 0.8) self.assertEqual( json.loads(extra_params["gpu_augmentations"]), [ diff --git a/research/vestibular_schwannoma/train_5fold.py b/research/vestibular_schwannoma/train_5fold.py new file mode 100644 index 0000000..b44c2dc --- /dev/null +++ b/research/vestibular_schwannoma/train_5fold.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python3 +"""Run reproducible five-fold VS training without using the notebook. + +This is the supported command-line entry point for cross-validation experiments. +Reusable training logic lives in ``workflow/`` and is shared with +``notebooks/01_five_fold_cross_validation.ipynb``; do not duplicate that logic here. + +Setup +----- +- Run in a fastMONAI environment with a CUDA GPU. SegMamba is optional. +- The default index is ``data/ml_dataset.csv``. Its image and mask paths expect + ``research/nii_data/`` beside ``research/vestibular_schwannoma/``. Pass another + index with ``--data-csv``; its relative paths resolve from the VS project directory. + +Usage (from ``research/vestibular_schwannoma``) +------------------------------------------------ +Inspect all options or run a short single-fold check:: + + python train_5fold.py --help + python train_5fold.py --models unet --folds 1 --epochs 5 --no-compile + +Run one model across all five folds:: + + python train_5fold.py --models unet + +Run the default five folds and skip an optional model if it is not installed:: + + python train_5fold.py --skip-unavailable + +The defaults request four models, folds 1-5, and 500 epochs. Models and folds run +sequentially so only one model occupies GPU memory at a time. Each held-out fold is +evaluated with TTA. + +Performance controls +-------------------- +- ``torch.compile`` is enabled for supported models. It costs startup time but can + speed long runs; use ``--no-compile`` for quick checks or compiler problems. +- ``--preprocess-workers N`` controls one-time preprocessing (default: up to 32). +- The patch queue defaults to 4 extraction workers and 300 buffered patches. If the + GPU waits for data and CPU/RAM are available, raise ``--queue-workers`` first and + then ``--queue-length``; a larger queue consumes more RAM. + +Preprocessing is cached in ``preprocessed/`` and outputs go below +``cv_results//`` unless overridden. Generated data, caches, MLflow +state, predictions, and weights stay outside Git through the project ``.gitignore``. + +This launcher performs cross-validation only. Use the shared workflow directly (or +extend the CLI explicitly) if an all-data final/deployment model is required. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from datetime import datetime, timezone +from pathlib import Path + +import pandas as pd +import torch + +from fastMONAI.vision_all import MedDataset, MedMask, ZNormalization, preprocess_dataset + + +PROJECT_ROOT = Path(__file__).resolve().parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from workflow.config import ExperimentConfig, make_patch_config # noqa: E402 +from workflow.models import ( # noqa: E402 + TRAINING_MODEL_CONFIGS, + get_training_model_configs, +) +from workflow.results import aggregate_results, build_model_comparison # noqa: E402 +from workflow.training import run_training_sweep # noqa: E402 + + +DEFAULT_MODELS = ("unet", "dynunet", "dynunet_small", "segmamba") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Train and evaluate the VS models with fixed five-fold cross-validation. " + "Models and folds run sequentially to bound GPU memory use." + ) + ) + parser.add_argument( + "--models", + nargs="+", + choices=tuple(TRAINING_MODEL_CONFIGS), + default=list(DEFAULT_MODELS), + help="Model keys to train in order (default: all four).", + ) + parser.add_argument( + "--folds", + nargs="+", + type=int, + default=[1, 2, 3, 4, 5], + help="Held-out folds to run (default: 1 2 3 4 5).", + ) + parser.add_argument("--epochs", type=int, default=500) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--learning-rate", type=float, default=1e-3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--foreground-probability", + type=float, + default=0.8, + help=( + "Probability of sampling a foreground-centred patch. The default 0.8 " + "preserves the original 80%% foreground / 20%% background control. " + "Pass 0.7 for the planned 70%% / 30%% comparison." + ), + ) + parser.add_argument( + "--samples-per-volume", + type=int, + default=4, + help="Patches extracted from each volume per epoch (default: 4).", + ) + parser.add_argument( + "--queue-workers", + type=int, + default=4, + help="CPU workers that fill the training patch queue (default: 4).", + ) + parser.add_argument( + "--queue-length", + type=int, + default=300, + help="Maximum patches buffered in RAM (default: 300).", + ) + parser.add_argument( + "--preprocess-workers", + type=int, + default=min(32, os.cpu_count() or 1), + help="CPU workers used for one-time preprocessing (default: up to 32).", + ) + parser.add_argument( + "--data-csv", + type=Path, + default=Path("data/ml_dataset.csv"), + help="Dataset index, relative to the VS project directory unless absolute.", + ) + parser.add_argument( + "--preprocessed-dir", + type=Path, + default=Path("preprocessed"), + help="Versioned preprocessing cache directory.", + ) + parser.add_argument( + "--results-root", + type=Path, + help="Output directory (default: cv_results/).", + ) + parser.add_argument( + "--no-compile", + action="store_true", + help="Disable torch.compile for models that support it.", + ) + parser.add_argument( + "--skip-unavailable", + action="store_true", + help="Skip an unavailable optional model such as SegMamba instead of failing early.", + ) + parser.add_argument( + "--stop-on-error", + action="store_true", + help="Stop immediately if one fold fails instead of continuing to the next run.", + ) + return parser + + +def _project_path(path: Path) -> Path: + return path if path.is_absolute() else PROJECT_ROOT / path + + +def _load_and_validate_dataset(data_csv: Path, folds: tuple[int, ...]) -> pd.DataFrame: + if not data_csv.is_file(): + raise FileNotFoundError(f"Dataset CSV not found: {data_csv}") + frame = pd.read_csv(data_csv) + required = {"case_id", "t1_img_path", "t1_seg_path", "fold"} + missing_columns = sorted(required - set(frame.columns)) + if missing_columns: + raise ValueError(f"Dataset CSV is missing columns: {missing_columns}") + if frame["case_id"].isna().any() or frame["case_id"].duplicated().any(): + raise ValueError("case_id values must be present and unique") + if frame[["t1_img_path", "t1_seg_path", "fold"]].isna().any().any(): + raise ValueError("Image paths, mask paths, and fold values must be present") + + available_folds = set(map(int, frame["fold"].unique())) + missing_folds = sorted(set(folds) - available_folds) + if missing_folds: + raise ValueError(f"Dataset does not contain requested folds: {missing_folds}") + + missing_files = [] + for column in ("t1_img_path", "t1_seg_path"): + for value in frame[column]: + path = Path(value) + if not path.is_file(): + missing_files.append(str(path)) + if len(missing_files) == 10: + break + if len(missing_files) == 10: + break + if missing_files: + shown = "\n ".join(missing_files) + raise FileNotFoundError(f"Missing dataset files (first 10):\n {shown}") + return frame + + +def _print_plan( + experiment: ExperimentConfig, + data_csv: Path, + results_root: Path, +) -> None: + print("\nFive-fold training plan") + print(f" Project: {PROJECT_ROOT}") + print(f" Dataset: {data_csv}") + print(f" Models: {', '.join(experiment.model_keys)}") + print(f" Folds: {experiment.folds}") + print(f" Epochs: {experiment.epochs}") + print(f" Batch size: {experiment.batch_size}") + print(f" Learning rate: {experiment.learning_rate}") + print( + " Patch sampling: " + f"{experiment.foreground_sampling_probability:.0%} foreground / " + f"{1 - experiment.foreground_sampling_probability:.0%} background" + ) + print(" Held-out inference TTA: ON (fixed)") + print(" Execution: sequential (one model/fold on the GPU at a time)") + print(f" Results: {results_root}\n") + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + os.chdir(PROJECT_ROOT) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA GPU is required for this training launcher") + + experiment = ExperimentConfig( + model_keys=tuple(args.models), + folds=tuple(args.folds), + run_cross_validation=True, + train_all_data=False, + training_seed=args.seed, + epochs=args.epochs, + batch_size=args.batch_size, + learning_rate=args.learning_rate, + use_tta=True, + compile_models=not args.no_compile, + target_spacing=(0.4102, 0.4102, 1.5), + patch_size=(192, 192, 48), + preprocess_workers=args.preprocess_workers, + samples_per_volume=args.samples_per_volume, + queue_num_workers=args.queue_workers, + queue_length=args.queue_length, + foreground_sampling_probability=args.foreground_probability, + continue_on_error=not args.stop_on_error, + ) + + data_csv = _project_path(args.data_csv) + preprocessed_dir = _project_path(args.preprocessed_dir) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + results_root = _project_path(args.results_root or Path("cv_results") / timestamp) + _print_plan(experiment, data_csv, results_root) + + model_configs = get_training_model_configs( + experiment.model_keys, + skip_unavailable=args.skip_unavailable, + ) + train_df = _load_and_validate_dataset(data_csv, experiment.folds) + print(f"Validated {len(train_df)} cases") + print(train_df["fold"].value_counts().sort_index().to_string()) + + normalization = [ZNormalization(masking_method="foreground")] + label_dataset = MedDataset( + img_list=train_df["t1_seg_path"].tolist(), + dtype=MedMask, + max_workers=experiment.preprocess_workers, + ) + dataset_version = label_dataset.fingerprint + if dataset_version is None: + raise RuntimeError("Could not fingerprint the segmentation masks") + + preprocessing_result = preprocess_dataset( + train_df, + img_col="t1_img_path", + mask_col="t1_seg_path", + output_dir=str(preprocessed_dir), + target_spacing=list(experiment.target_spacing), + apply_reorder=True, + transforms=normalization, + max_workers=experiment.preprocess_workers, + dataset_version=dataset_version, + ) + print(f"Dataset version: {dataset_version}") + print(f"Preprocessing cache: {preprocessing_result.cache_version}") + print(f"Cache reused: {preprocessing_result.reused}") + print(f"Manifest: {preprocessing_result.manifest_path}") + + torch.backends.cudnn.benchmark = True + sweep = run_training_sweep( + model_configs, + train_df, + experiment=experiment, + patch_config=make_patch_config(experiment, normalization), + preprocessing_manifest=preprocessing_result.manifest_path, + results_root=results_root, + ) + + combined = {} + for model_key in model_configs: + result = aggregate_results( + results_root / model_key, + experiment.folds, + train_df, + ) + if result is not None: + combined[model_key] = result + comparison = build_model_comparison(combined) + if not comparison.empty: + comparison_path = results_root / "cv_model_comparison.csv" + comparison.to_csv(comparison_path) + print(f"Cross-model comparison: {comparison_path}") + + if sweep.failures: + print(f"Training finished with {len(sweep.failures)} failed run(s):") + for failure in sweep.failures: + location = f" fold {failure.fold}" if failure.fold is not None else "" + print( + f" {failure.model_key}{location}: " + f"{failure.error_type}: {failure.message}" + ) + return 1 + + print(f"All requested runs completed: {results_root}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/research/vestibular_schwannoma/workflow/config.py b/research/vestibular_schwannoma/workflow/config.py index 59a1036..821eeef 100644 --- a/research/vestibular_schwannoma/workflow/config.py +++ b/research/vestibular_schwannoma/workflow/config.py @@ -38,6 +38,7 @@ class ExperimentConfig: samples_per_volume: int = 4 queue_num_workers: int = 4 queue_length: int = 300 + foreground_sampling_probability: float = 0.8 continue_on_error: bool = True def __post_init__(self) -> None: @@ -81,6 +82,10 @@ def __post_init__(self) -> None: raise ValueError("target_spacing must contain three positive values") if len(self.patch_size) != 3 or any(value <= 0 for value in self.patch_size): raise ValueError("patch_size must contain three positive values") + if not 0 < self.foreground_sampling_probability < 1: + raise ValueError( + "foreground_sampling_probability must be strictly between 0 and 1" + ) def make_patch_config(config: ExperimentConfig, normalization: list) -> PatchConfig: @@ -90,7 +95,10 @@ def make_patch_config(config: ExperimentConfig, normalization: list) -> PatchCon patch_size=list(config.patch_size), samples_per_volume=config.samples_per_volume, sampler_type="label", - label_probabilities={0: 0.2, 1: 0.8}, + label_probabilities={ + 0: round(1 - config.foreground_sampling_probability, 12), + 1: config.foreground_sampling_probability, + }, patch_overlap=0.5, keep_largest_component=False, target_spacing=list(config.target_spacing), diff --git a/research/vestibular_schwannoma/workflow/models.py b/research/vestibular_schwannoma/workflow/models.py index fe09cc4..9014f87 100644 --- a/research/vestibular_schwannoma/workflow/models.py +++ b/research/vestibular_schwannoma/workflow/models.py @@ -80,25 +80,32 @@ def _make_dynunet_loss() -> CustomLoss: }, ) -DYNUNET_SPEC = make_model_spec( - "monai.dynunet", - { - "spatial_dims": 3, - "in_channels": 1, - "out_channels": 2, - "kernel_size": [[3, 3, 3]] * 5, - "strides": [[1, 1, 1]] + [[2, 2, 2]] * 4, - "upsample_kernel_size": [[2, 2, 2]] * 4, - "filters": [64, 128, 256, 512, 1024], - "res_block": True, - "deep_supervision": True, - "deep_supr_num": 3, - }, - wrapper_spec={ - "wrapper_id": "fastmonai.dynunet_ds_adapter", - "wrapper_kwargs": {}, - }, -) + +def _make_dynunet_spec(filters: list[int]) -> dict: + return make_model_spec( + "monai.dynunet", + { + "spatial_dims": 3, + "in_channels": 1, + "out_channels": 2, + "kernel_size": [[3, 3, 3]] * 5, + "strides": [[1, 1, 1]] + [[2, 2, 2]] * 4, + "upsample_kernel_size": [[2, 2, 2]] * 4, + "filters": filters, + "res_block": True, + "deep_supervision": True, + "deep_supr_num": 3, + }, + wrapper_spec={ + "wrapper_id": "fastmonai.dynunet_ds_adapter", + "wrapper_kwargs": {}, + }, + ) + + +DYNUNET_SPEC = _make_dynunet_spec([64, 128, 256, 512, 1024]) + +DYNUNET_SMALL_SPEC = _make_dynunet_spec([32, 64, 128, 256, 512]) SEGMAMBA_SPEC = make_model_spec( "segmamba.v2", @@ -128,6 +135,13 @@ def _make_dynunet_loss() -> CustomLoss: make_loss=_make_dynunet_loss, experiment_name="vestibular_schwannoma_dynunet", ), + "dynunet_small": TrainingModelConfig( + key="dynunet_small", + display_name="DynUNet Small (32-512)", + model_spec=DYNUNET_SMALL_SPEC, + make_loss=_make_dynunet_loss, + experiment_name="vestibular_schwannoma_dynunet_small", + ), "segmamba": TrainingModelConfig( key="segmamba", display_name="SegMamba V2", diff --git a/research/vestibular_schwannoma/workflow/training.py b/research/vestibular_schwannoma/workflow/training.py index 1303d89..e5a91e7 100644 --- a/research/vestibular_schwannoma/workflow/training.py +++ b/research/vestibular_schwannoma/workflow/training.py @@ -282,6 +282,9 @@ def train_one_fold( extra_tags={"fold": str(fold), "run_group": Path(results_dir).parent.name}, extra_params={ "training_seed": experiment.training_seed, + "foreground_sampling_probability": ( + experiment.foreground_sampling_probability + ), "gpu_augmentations": _gpu_augmentations_json(gpu_augmentation), }, preprocessing_manifest=preprocessing_manifest, @@ -379,6 +382,9 @@ def train_all_data_model( }, extra_params={ "training_seed": experiment.training_seed, + "foreground_sampling_probability": ( + experiment.foreground_sampling_probability + ), "gpu_augmentations": _gpu_augmentations_json(gpu_augmentation), }, preprocessing_manifest=preprocessing_manifest,