# claim-berlin/mbh-seg-2024-winning-solution architecture
> Architecture diagram of the GitHub repository claim-berlin/mbh-seg-2024-winning-solution, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-24). It is an AI-made overview: check details against the source.
- Interactive diagram: https://gitdiagram.com/claim-berlin/mbh-seg-2024-winning-solution
- Repository: https://github.com/claim-berlin/mbh-seg-2024-winning-solution
## Overview
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: [`nnunetv2/training/nnUNetTrainer`](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/nnUNetTrainer) (Trainer Core), [`nnunetv2/inference`](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/inference) (Inference), [`submission.py`](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/submission.py) (Submission), [`nnunetv2/preprocessing`](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/preprocessing) (Preprocess)
## Diagram (Mermaid)
```mermaid
flowchart TD
subgraph group_entry["Deployment"]
node_submission["Submission
entrypoint
[submission.py]"]
end
subgraph group_data["Data Prep"]
node_aux_inputs["Aux Inputs
prelabeling input"]
node_imageio["Image I/O
reader-writer"]
node_preprocessing["Preprocess
array prep"]
end
subgraph group_core["nnU-Net Core"]
node_plan["Plan & Fingerprint
dataset planning"]
node_trainer["Trainer Core
training core"]
node_dataloading["Data Loaders
batch loading"]
node_augment["Augmentations
transforms"]
node_losses["Losses
loss family"]
node_predict["Inference
sliding-window"]
node_weights["Weights"]
node_postprocess["Postprocess
cleanup"]
end
subgraph group_ops["Supporting Ops"]
node_sharing["Model Sharing
artifact mgmt"]
node_ensemble["Ensembling
prediction merge
[ensemble.py]"]
node_eval["Evaluation
validation"]
node_convert["Dataset Convert
conversion scripts"]
node_batchrun["Batch Runs
offline ops"]
node_tests["Integration Tests
validation suite"]
end
node_submission -->|"reads/writes"| node_imageio
node_submission -->|"consumes"| node_aux_inputs
node_submission -->|"prepares"| node_preprocessing
node_preprocessing -->|"uses backends"| node_imageio
node_preprocessing -->|"follows plans"| node_plan
node_plan -->|"configures"| node_trainer
node_plan -.->|"defines"| node_preprocessing
node_trainer -->|"pulls batches"| node_dataloading
node_trainer -->|"applies"| node_augment
node_trainer -->|"optimizes"| node_losses
node_trainer -->|"loads"| node_weights
node_predict -->|"exports via"| node_imageio
node_predict -->|"runs with"| node_weights
node_predict -->|"hands off"| node_postprocess
node_submission -->|"invokes"| node_predict
node_ensemble -->|"merges outputs"| node_predict
node_sharing -->|"distributes"| node_weights
node_eval -->|"scores"| node_trainer
node_convert -->|"feeds datasets"| node_plan
node_batchrun -.->|"collects"| node_eval
node_tests -.->|"verifies"| node_submission
node_tests -.->|"checks"| node_predict
click node_submission "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/submission.py"
click node_imageio "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/imageio"
click node_preprocessing "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/preprocessing"
click node_plan "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/experiment_planning"
click node_trainer "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/nnUNetTrainer"
click node_dataloading "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/dataloading"
click node_augment "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/data_augmentation"
click node_losses "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/loss"
click node_predict "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/inference"
click node_weights "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/run/load_pretrained_weights.py"
click node_postprocess "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/postprocessing/remove_connected_components.py"
click node_sharing "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/model_sharing"
click node_ensemble "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/ensembling/ensemble.py"
click node_eval "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/evaluation"
click node_convert "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/dataset_conversion"
click node_batchrun "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/batch_running"
click node_tests "https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/tests/integration_tests"
classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_submission toneBlue
class node_aux_inputs,node_imageio,node_preprocessing toneAmber
class node_plan,node_trainer,node_dataloading,node_augment,node_losses,node_predict,node_weights,node_postprocess toneMint
class node_sharing,node_ensemble,node_eval,node_convert,node_batchrun,node_tests toneRose
```
## Components
### Deployment
Challenge entrypoint and packaged inference boundary.
- [Submission](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/submission.py): entrypoint · Reads study folders and writes final segmentations. · `submission.py`
### Data Prep
I/O, preprocessing, and auxiliary input handling.
- **Aux Inputs**: prelabeling input · Upstream TotalSegmentator and other model predictions used as extra channels.
- [Image I/O](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/imageio): reader-writer · Backend registry for NIfTI, SimpleITK, TIFF, and natural-image formats. · `nnunetv2/imageio`
- [Preprocess](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/preprocessing): array prep · Cropping, normalization, and resampling before inference or training. · `nnunetv2/preprocessing`
### nnU-Net Core
Planning, training, and inference engine.
- [Plan & Fingerprint](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/experiment_planning): dataset planning · Computes fingerprints, plans target spacing, patch sizes, and topology. · `nnunetv2/experiment_planning`
- [Trainer Core](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/nnUNetTrainer): training core · Main nnU-Net trainer abstraction and the trainer family variants. · `nnunetv2/training/nnUNetTrainer`
- [Data Loaders](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/dataloading): batch loading · 2D/3D loaders and skeleton-aware sampling for training batches. · `nnunetv2/training/dataloading`
- [Augmentations](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/data_augmentation): transforms · Spatial and task-specific transforms injected during training. · `nnunetv2/training/data_augmentation`
- [Losses](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/training/loss): loss family · Standard Dice/CE plus boundary and skeleton-aware losses. · `nnunetv2/training/loss`
- [Inference](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/inference): sliding-window · Patch-wise prediction over large volumes and export to disk. · `nnunetv2/inference`
- [Weights](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/run/load_pretrained_weights.py): pretrained weights · Loads packaged checkpoints and connects them to the network. · `nnunetv2/run/load_pretrained_weights.py`
- [Postprocess](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/postprocessing/remove_connected_components.py): cleanup · Challenge-specific connected-component cleanup and label reassignment. · `nnunetv2/postprocessing/remove_connected_components.py`
### Supporting Ops
Packaging, evaluation, conversion, and validation utilities.
- [Model Sharing](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/model_sharing): artifact mgmt · Import, export, and download entry points for distributed models. · `nnunetv2/model_sharing`
- [Ensembling](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/blob/main/nnunetv2/ensembling/ensemble.py): prediction merge · Combines outputs from multiple models when used. · `nnunetv2/ensembling/ensemble.py`
- [Evaluation](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/evaluation): validation · Cross-validation aggregation and best-configuration selection. · `nnunetv2/evaluation`
- [Dataset Convert](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/dataset_conversion): conversion scripts · Dataset-specific converters and integration-test dataset builders. · `nnunetv2/dataset_conversion`
- [Batch Runs](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/batch_running): offline ops · Cluster benchmarking, result collection, and release-training helpers. · `nnunetv2/batch_running`
- [Integration Tests](https://github.com/claim-berlin/mbh-seg-2024-winning-solution/tree/main/nnunetv2/tests/integration_tests): validation suite · End-to-end training and inference checks for reproducibility. · `nnunetv2/tests/integration_tests`
## Connections
- 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.