$ show MIC-DKFZ/nnUNet
nnU-Net v2 is a self-configuring deep learning framework for medical image segmentation that automates preprocessing, network architecture selection, training, and inference. It ingests raw medical images in various formats (NIfTI, TIFF, DICOM via SimpleITK) and produces segmentation masks. The system is organized as a modular monolith: experiment planning extracts dataset fingerprints to determine optimal preprocessing and network topology; a preprocessing pipeline handles resampling, normalization, and cropping; training uses a configurable base trainer with many variants (including 3D/2D U-Nets, residual encoders, and transformer-based Primus architectures); inference employs sliding-window with Gaussian weighting and test-time augmentation. Key entry points are CLI commands (nnUNetv2_plan_and_preprocess, nnUNetv2_train, nnUNetv2_predict) and the programmatic nnUNetPredictor. The codebase demonstrates high craftsmanship in its checkpoint system (safetensors with backward compatibility), pluggable image I/O, and configuration management via plans files, though it carries some technical debt from legacy .pth support and a proliferation of trainer variant subclasses.
The LLM identified 5 files worth highlighting for reuse:
Provides PlansManager and ConfigurationManager classes that parse and manage the experiment plans JSON file. Supports inheritance (inherits_from) to layer configurations, caching, and resolution of resampling functions. This is a clean, extensible configuration system that can be extracted for other projects needing hierarchical JSON-based configuration.
Implements a robust checkpoint storage system using safetensors format with a metadata header. Supports atomic writes, readback verification, legacy .pth fallback, partial loading (keys parameter), and conversion tools. Demonstrates high craftsmanship in handling format migration and backward compatibility, and is safely reusable in any PyTorch project.
Defines BaseReaderWriter abstract class and a registry (reader_writer_registry.py) for pluggable image I/O. Concrete implementations support NIfTI (nibabel), NRRD/MHA (SimpleITK), TIFF, and natural images. The pattern is clean, minimal, and easily extracted for any project needing format-agnostic data loading.
Implements ForegroundLocationsWriter and ForegroundLocations for compressed storage of foreground voxel coordinates. Uses a CSR-like sparse index with optional Blosc2 compression, dramatically reducing disk space compared to legacy pickle files. Includes backward-compatible loading and a command-line extraction tool. This is a self-contained component for efficient random access to foreground sampling locations.
Provides torch-based resampling functions (resample_torch_simple, resample_torch_fornnunet) for medical images and segmentation masks. Uses efficient gather operations for out-of-plane resampling, running argmax for segmentation with tie-breaking, and separate-z handling. Validated against scipy-based reference for bit-exact correctness. Demonstrates non-obvious optimization for interpolation of 3D medical volumes.