How claim-berlin/mbh-seg-2024-winning-solution works
1. Repository role
- This is an nnU-Net v2–based segmentation codebase, adapted for the MBHSEG challenge solution.
- The practical deployment boundary is submission.py, which takes a folder of medical images and writes segmentation outputs.
2. Top-level inference boundary
- The expected runtime input is a directory of NIfTI studies.
- The pipeline produces multiclass brain hemorrhage masks, and the README indicates two model bundles are expected under models/multiclass and models/prelabeling.
3. External inputs used by the model
- The solution is not image-only: it uses TotalSegmentator-derived segmentations and other open-source model outputs as an additional input channel.
- Architecturally, these are upstream auxiliary predictions that feed into the nnU-Net inference pipeline rather than separate post-hoc heuristics.
4. Image I/O and modality handling
- The imageio package contains reader/writer backends for NIfTI via nibabel and SimpleITK, plus TIFF and natural-image support.
- This layer isolates file-format concerns from preprocessing and inference, so the core pipeline operates on standardized array/tensor representations.
5. Preprocessing pipeline
- preprocessing handles cropping, normalization, and resampling before inference/training.
- The resampling code is split into default CPU-style logic and torch-based variants, which matters for how spatial harmonization is performed before the network sees the data.
6. Planning and dataset fingerprinting
- experiment_planning computes dataset fingerprints and plans configurations such as target spacing, patch sizes, and network topology.
- This is the subsystem that turns raw dataset properties into a trainable/inferable nnU-Net plan, and it also contains integrity checks and pretraining plan transfer utilities.
7. Model and trainer core
- The main training abstraction is training/nnUNetTrainer/nnUNetTrainer.py, with many variants for augmentation, loss, optimizer, sampling, and training length.
- The README’s winning setup corresponds to the ResEncL family with the DA5 augmentation-heavy trainer, while keeping standard nnU-Net losses and hyperparameters.
8. Training data flow
- training/dataloading provides 2D/3D loaders and skeleton-aware variants, while training/data_augmentation injects spatial and task-specific transforms.
- Losses include standard Dice/CE variants plus specialized boundary/skeleton losses, showing that training is modular and configuration-driven rather than monolithic.
9. Inference execution
- inference/predict_from_raw_data.py and sliding_window_prediction.py implement the core forward pass over large volumes.
- export_prediction.py and the iterator utilities manage batching, patch-wise inference, and writing results back to disk; run/load_pretrained_weights.py connects the packaged weights to the network.
10. Postprocessing and cleanup
- postprocessing/remove_connected_components.py implements the challenge-specific cleanup described in the README.
- The key rule is to resolve connected components containing both subdural and epidural labels by reassignment to the dominant class, and to remove very small SAH/epidural/subdural components below the voxel threshold.
11. Model packaging and sharing
- model_sharing contains import/export/download entry points for moving trained weights between environments.
- This is important for deployment because the repository assumes model artifacts are distributed separately from source code.
12. Evaluation, ensembling, and selection
- evaluation includes cross-validation aggregation and best-configuration selection, while ensembling/ensemble.py combines predictions when multiple models are used.
- These modules are primarily offline/validation-side, but they explain how the final checkpoint/configuration was chosen and combined.
13. Dataset conversion and experimental utilities
- dataset_conversion contains many dataset-specific conversion scripts, including integration-test datasets and decathlon converters.
- batch_running and the tests/integration_tests folder support cluster execution, benchmark runs, and end-to-end validation, indicating the codebase was used both for research training and controlled reproducibility checks.
14. Overall architecture summary
- The core flow is: raw CT data + auxiliary segmentation input → I/O and preprocessing → nnU-Net plan/trainer → sliding-window inference → challenge-specific postprocessing → exported mask.
- Everything else in the repo supports one of three concerns: preparing datasets, training/validating models, or packaging/invoking the final inference workflow.
Main entry points: nnUNetTrainer Trainer Core · inference Inference · submission.py Submission · preprocessing Preprocess
Components 18
Deployment
Challenge entrypoint and packaged inference boundary.
- Submission · Reads study folders and writes final segmentations.
submission.py
Data Prep
I/O, preprocessing, and auxiliary input handling.
- Aux Inputs · Upstream TotalSegmentator and other model predictions used as extra channels.
- Image I/O · Backend registry for NIfTI, SimpleITK, TIFF, and natural-image formats.
nnunetv2/imageio - Preprocess · Cropping, normalization, and resampling before inference or training.
nnunetv2/preprocessing
nnU-Net Core
Planning, training, and inference engine.
- Plan & Fingerprint · Computes fingerprints, plans target spacing, patch sizes, and topology.
nnunetv2/experiment_planning - Trainer Core · Main nnU-Net trainer abstraction and the trainer family variants.
nnunetv2/training/nnUNetTrainer - Data Loaders · 2D/3D loaders and skeleton-aware sampling for training batches.
nnunetv2/training/dataloading - Augmentations · Spatial and task-specific transforms injected during training.
nnunetv2/training/data_augmentation - Losses · Standard Dice/CE plus boundary and skeleton-aware losses.
nnunetv2/training/loss - Inference · Patch-wise prediction over large volumes and export to disk.
nnunetv2/inference - Weights · Loads packaged checkpoints and connects them to the network.
nnunetv2/run/load_pretrained_weights.py - Postprocess · Challenge-specific connected-component cleanup and label reassignment.
nnunetv2/postprocessing/remove_connected_components.py
Supporting Ops
Packaging, evaluation, conversion, and validation utilities.
- Model Sharing · Import, export, and download entry points for distributed models.
nnunetv2/model_sharing - Ensembling · Combines outputs from multiple models when used.
nnunetv2/ensembling/ensemble.py - Evaluation · Cross-validation aggregation and best-configuration selection.
nnunetv2/evaluation - Dataset Convert · Dataset-specific converters and integration-test dataset builders.
nnunetv2/dataset_conversion - Batch Runs · Cluster benchmarking, result collection, and release-training helpers.
nnunetv2/batch_running - Integration Tests · End-to-end training and inference checks for reproducibility.
nnunetv2/tests/integration_tests
Connections 22
- Submission → Image I/O: reads/writes · Loads NIfTI studies and writes predicted masks.
- Submission → Aux Inputs: consumes · Uses auxiliary segmentation channels alongside CT data.
- Submission → Preprocess: prepares · Standardizes inputs before model inference.
- Preprocess → Image I/O: uses backends · Relies on format-specific readers and writers.
- Preprocess → Plan & Fingerprint: follows plans · Uses spacing, patch, and normalization decisions from planning.
- Plan & Fingerprint → Trainer Core: configures · Provides dataset-derived training and inference configuration.
- Plan & Fingerprint → Preprocess: defines · Determines resampling and preprocessing behavior.
- Trainer Core → Data Loaders: pulls batches · Consumes training samples through the loader layer.
- Trainer Core → Augmentations: applies · Injects augmentation pipelines during training.
- Trainer Core → Losses: optimizes · Uses selected losses and trainer variants.
- Trainer Core → Weights: loads · Initializes or restores model parameters.
- Inference → Image I/O: exports via · Writes final predictions using the I/O layer.
- Inference → Weights: runs with · Uses pretrained checkpoints for inference.
- Inference → Postprocess: hands off · Sends raw predictions to cleanup and label correction.
- Submission → Inference: invokes · Triggers the end-to-end inference pipeline.
- Ensembling → Inference: merges outputs · Combines multiple prediction streams when configured.
- Model Sharing → Weights: distributes · Moves model artifacts between environments.
- Evaluation → Trainer Core: scores · Evaluates trained configurations and CV results.
- Dataset Convert → Plan & Fingerprint: feeds datasets · Produces standardized datasets for planning and training.
- Batch Runs → Evaluation: collects · Aggregates offline benchmarking and training results.
- Integration Tests → Submission: verifies · Exercises the deployment boundary end to end.
- Integration Tests → Inference: checks · Validates inference behavior and configuration selection.