How mic-dkfz/nnlandmark works
1. What this repository is
nnLandmark is a fork of nnU-Net specialized for 3D medical landmark detection. The core architectural idea is to reuse nnU-Net’s self-configuring pipeline, but represent landmarks as multi-label segmentation targets and then recover landmark coordinates from predicted heatmaps/segmentations.
2. Top-level pipeline
The main data flow is: raw medical volumes and landmark annotations in nnU-Net-style folders and JSONs → dataset conversion / validation → experiment planning and preprocessing → trainer-based model fitting → inference on raw cases → coordinate extraction and JSON export → evaluation in voxel and millimeter space.
3. Repository boundaries
The code is split into the same major nnU-Net-style layers:
- dataset_conversion: dataset-specific importers and split generation
- experiment_planning: fingerprints, plans, and preprocessing decisions
- preprocessing: cropping, resampling, normalization
- training: trainers, losses, augmentations, dataloading
- inference: sliding-window prediction and export
- evaluation: landmark metrics and dataset-specific scoring
- utilities: shared plumbing for plans, labels, paths, class lookup, and JSON handling
4. Data model and storage convention
The repo uses nnU-Net’s environment-variable path model, but under nnLM_raw, nnLM_preprocessed, and nnLM_results. Training data lives in imagesTr and labelsTr. Unlike class segmentation, each landmark class is encoded as a label value in a multi-label segmentation map, often as a small 3×3×3 cube around the target voxel. Additional JSON files hold landmark coordinates and spacing metadata.
5. Dataset conversion layer
nnlandmark/dataset_conversion contains both generic converters and many dataset-specific scripts for public landmark datasets such as AFIDs, MML, DMGLD LFC, PDDCA, FeTA, and several challenge/benchmark datasets. These scripts standardize raw data into the nnU-Net-style layout, create dataset metadata, and define official or custom train/test splits.
6. Experiment planning and preprocessing
The experiment_planning module mirrors nnU-Net’s self-configuration logic. It extracts dataset fingerprints, verifies integrity, chooses resampling and patch/layout decisions, and writes plans used later by training and inference. There are explicit planner variants for residual encoder U-Nets and resampling/no-resampling options.
7. Preprocessing mechanics
Preprocessing is composed of cropping, normalization, and resampling modules. This is where image statistics and spacing assumptions are reconciled before training. The design keeps modality-specific normalization and resampling separate, so dataset-specific behavior can be adjusted without changing the rest of the pipeline.
8. Training subsystem
Training is centered on nnUNetTrainer and a set of variants. The base trainer is extended by:
- generic variants for loss, augmentation, optimizer, sampling, and training length
- project-specific trainers for nnLandmark
- Kaggle/competition-specific trainers and losses
The training stack includes custom data augmentation, deep supervision, Dice/CE/top-k/regression-style losses, and learning-rate schedulers such as polynomial decay and warmup.
9. Project-specific model branches
There are two notable specialization paths:
- project_specific/nnLandmark: landmark-oriented trainer and architectures, including a BiFormer-based U-Net variant
- project_specific/kaggle2025_byu: a separate branch for motor/coordinate regression and competition-specific losses, coordinate representations, and augmentation
This indicates the repo supports both segmentation-style landmark detection and some regression-like competition workflows.
10. Inference and coordinate extraction
Inference is handled by inference/predict_from_raw_data.py and sliding-window prediction utilities. The output is not only a segmentation mask but also per-case prediction JSONs containing voxel coordinates and confidence-like values for each landmark. The README describes a postprocessing step that turns heatmap predictions into multi-label segmentations by taking the top 27 voxels per landmark.
11. Evaluation and analysis
evaluation computes landmark error metrics and produces summary reports in both voxel and millimeter units. There are dataset-specific evaluators for nnLandmark benchmarks and a separate FeTA biometry measurement evaluator that measures distances between paired landmarks using spacing metadata. This module is the boundary where predicted coordinates are converted into final challenge/reporting metrics.
12. Supporting infrastructure
- imageio abstracts I/O across NIfTI, SimpleITK, TIFF, and natural-image readers/writers.
- ensembling combines model outputs.
- postprocessing contains cleanup such as connected-component filtering.
- model_sharing supports export/import/download of trained models.
- run contains training launch helpers and pretrained-weight loading.
These modules keep model code, I/O, and deployment concerns separated.
13. Batch execution and large-scale runs
batch_running and the tests/integration_tests area show that the repo is designed for cluster-style execution, including LSF job generation, benchmarking, result collection, and release-training scripts. This matters architecturally because many workflows are intended to run at scale rather than interactively.
14. Quality and reproducibility
The repository includes GitHub Actions for spell checking and tests, plus integration-test datasets and scripts. That suggests a strong focus on reproducible end-to-end pipeline behavior across dataset conversion, planning, training, and inference.
15. How to read the architecture diagram
The central hub is the nnU-Net-like training/inference pipeline. Data enters through dataset conversion and standardized JSON/folder formats, flows through planning and preprocessing, then into configurable trainers. Outputs return through inference, export, and evaluation, with dataset-specific scripts and competition utilities branching off around the edges.
Main entry points: nnUNetTrainer.py Base trainer · dataset_conversion Dataset conversion · integration_tests Integration tests · experiment_planning Planning
Components 19
Data & setup
Dataset conversion, fingerprints, preprocessing, and shared file/label conventions.
- Raw data · Source volumes and landmark annotations in nnU-Net-style folders and metadata JSONs.
- Dataset conversion · Dataset-specific converters that standardize public landmark datasets into nnU-Net layout and define splits.
nnlandmark/dataset_conversion - Dataset JSONs · Dataset description files and split definitions used by planning, training, and evaluation.
nnlandmark/dataset_conversion/generate_dataset_json.py - Planning · Extracts dataset fingerprints, verifies integrity, and writes configuration plans for preprocessing and training.
nnlandmark/experiment_planning - Preprocessing · Crops, rescales, and normalizes data before training using modality- and spacing-aware logic.
nnlandmark/preprocessing - Path model · Environment-variable driven roots for raw, preprocessed, and results storage.
nnlandmark/paths.py
Training core
Model training, losses, augmentation, dataloading, and trainer variants.
- Base trainer · Core training loop and nnU-Net training conventions reused across all project variants.
nnlandmark/training/nnUNetTrainer/nnUNetTrainer.py - Trainer variants · Reusable training variants for augmentation, loss functions, optimizers, sampling, schedules, and training length.
nnlandmark/training/nnUNetTrainer/variants - nnLandmark trainer · Project-specific landmark detection trainer and architectures, including the BiFormer U-Net branch.
nnlandmark/training/nnUNetTrainer/project_specific/nnLandmark/nnLandmark_trainer.py - Kaggle branch · Competition-specific trainer stack with coordinate representations, custom losses, and extra augmentation.
nnlandmark/training/nnUNetTrainer/project_specific/kaggle2025_byu/MotorRegressionTrainer.py - Data & loss · Dataloaders, augmentation transforms, loss implementations, and LR schedulers feeding the trainer.
nnlandmark/training
Inference & eval
Prediction, coordinate extraction, ensembling, postprocessing, and evaluation.
- Raw inference · Loads raw cases, runs sliding-window prediction, and writes per-case outputs.
nnlandmark/inference/predict_from_raw_data.py - Coord export · Converts predictions into landmark coordinates and JSON exports for downstream use.
nnlandmark/inference/export_prediction.py - Postprocess · Combines model outputs and optionally cleans masks before coordinate recovery.
nnlandmark/ensembling/ensemble.py - Evaluation · Computes voxel/mm errors, dataset-specific scores, and summary reports.
nnlandmark/evaluation
Ops & support
Batch execution, model sharing, utilities, and test scaffolding.
- Batch runs · LSF and benchmarking scripts for large-scale training, inference, and result collection.
nnlandmark/batch_running - Model sharing · Utilities for packaging, downloading, and importing trained models.
nnlandmark/model_sharing - Utilities · Shared helpers for plans, labels, JSON export, path handling, and class lookup.
nnlandmark/utilities - Integration tests · End-to-end test scripts and example data validating conversion, training, and inference flows.
nnlandmark/tests/integration_tests
Connections 22
- Raw data → Dataset conversion: import · Standardize raw datasets into nnU-Net-style landmark training cases.
- Dataset conversion → Dataset JSONs: write metadata · Generate dataset descriptions, label maps, and split definitions.
- Dataset JSONs → Planning: plan from · Use dataset metadata and splits to derive fingerprints and plans.
- Planning → Preprocessing: configure · Planning decisions determine cropping, resampling, and normalization behavior.
- Path model → Dataset conversion: root dirs · Provides shared raw/preprocessed/results locations for converters and pipelines.
- Path model → Base trainer: results dirs · Training resolves output locations through the shared path model.
- Preprocessing → Base trainer: feed data · Preprocessed volumes and labels become training inputs.
- Data & loss → Base trainer: used by · Augmentation, loaders, losses, and schedulers are assembled inside trainers.
- Base trainer → Trainer variants: specialize · Variants change augmentation, optimization, sampling, architecture, or training length.
- Base trainer → nnLandmark trainer: specialize · Project-specific landmark trainer extends the base nnU-Net loop.
- Base trainer → Kaggle branch: specialize · Competition branch adapts the base trainer for motor/coordinate regression workflows.
- Base trainer → Model sharing: publish · Trained models can be exported, imported, or downloaded for reuse.
- Base trainer → Raw inference: run inference · Trained checkpoints are applied to raw cases through the inference pipeline.
- Raw inference → Postprocess: predict · Prediction outputs may be ensembled and cleaned before export.
- Postprocess → Coord export: convert · Masks or heatmaps are reduced to landmark coordinates and per-case JSON.
- Coord export → Evaluation: score · Exported coordinates are measured in voxel and millimeter space.
- Evaluation → Batch runs: aggregate · Batch scripts collect results and summarize runs across datasets and jobs.
- Utilities → Planning: support · Plans, labels, and helper utilities underpin planning and configuration lookup.
- Utilities → Coord export: support · Shared JSON and collation helpers support export formatting.
- Integration tests → Dataset conversion: exercise · Integration tests cover dataset conversion and standardized layouts.
- Integration tests → Raw inference: exercise · Integration tests validate end-to-end inference behavior.
- Integration tests → Base trainer: exercise · Integration tests cover training and configuration flows.