MIR: Medical Image Registration toolkit
MIR is a PyTorch research toolkit for deformable medical image registration. It brings
state-of-the-art registration networks (TransMorph, VFA, SITReg, VoxelMorph, LessNet, Fourier-Net and
their variants), classical baselines, training losses, evaluation metrics and reproducible benchmark
pipelines together behind one interface: a (moving, fixed) pair in, a dense displacement field out.
Documentation: https://junyuchen.me/MIR
Highlights
- One convention for every model. A registration network takes a
(moving, fixed) pair of
[B, 1, D, H, W] tensors and returns a displacement field in voxels that SpatialTransformer applies
to any image or label map. Swapping architectures is a one-line change.
- Diffeomorphic by design. Choose the transformation model: displacement, stationary or time-varying
velocity fields, symmetric inverse-consistent mappings, or LDDMM geodesic shooting.
- Spatially varying regularization (SPR): learn where the deformation should be smooth and where it
may be sharp, with a per-voxel regularization map and a Beta prior.
- Hypernetwork conditioning (HyperMorph, HyperTransMorph, HyperVFA) to cover a range of
regularization strengths with one model.
- Mixed precision and multi-GPU. Train and run any model in float16 or bfloat16 and on several GPUs with
DistributedDataParallel or DataParallel; losses and warping stay in float32 for accuracy
(docs, example).
- Runs in the browser. Models export to ONNX; the
registration playground runs a 2D HyperVFA-SPR with ONNX Runtime Web
(tutorial, docs).
- Classical and third-party baselines behind the same interface: ConvexAdam, deedsBCV, FireANTs,
SITReg, MultiMorph, intensity normalization.
- Training and evaluation batteries: NCC / SSIM / MIND / mutual information / correlation ratio /
Dice losses, diffusion and local regularizers, Beta priors, GradICON, Jacobian-based folding metrics,
augmentation, and a registry of pretrained weights.
Installation
Install a CUDA-enabled PyTorch build first, then MIR in editable mode:
git clone https://github.com/junyuchen245/MIR.git
cd MIR
python3.8 -m pip install -U pip
python3.8 -m pip install -e .
Python 3.8+ is required. deedsBCV is Linux-only; everything else is platform independent.
See the installation guide for optional
dependencies and the pretrained weights.
Sixty-second example
import torch, nibabel as nib
from MIR import ModelWeights
from MIR.models import VFA, SpatialTransformer
import MIR.models.configs_VFA as configs_VFA
mov = torch.from_numpy(nib.load('subject.nii.gz').get_fdata()[None, None]).float().cuda()
fix = torch.from_numpy(nib.load('template.nii.gz').get_fdata()[None, None]).float().cuda()
config = configs_VFA.get_VFA_default_config()
config.img_size = fix.shape[2:]
model = VFA(config, device='cuda:0').cuda().eval()
model.load_state_dict(torch.load('VFA_LUMIR24.pth')[ModelWeights['VFA-LUMIR24-MonoModal']['wts_key']])
with torch.no_grad():
flow = model((mov, fix)) # [1, 3, D, H, W] displacement in voxels
warped = SpatialTransformer(fix.shape[2:]).cuda()(mov, flow)
Models
| Family | Summary | Docs |
|---|
| TransMorph, TransMorph-TVF, TransMorph-TVF-SPR, TransMorph-DCA | Swin-Transformer registration networks | TransMorph |
| VFA, VFA-SPR, TransVFA, HyperVFA | Vector field attention | VFA |
| HyperTransMorph, HyperMorph | Hypernetwork-conditioned regularization | HyperTransMorph |
| VoxelMorph, LessNet, Fourier-Net(+) | Convolutional baselines (with SVF and SPR variants) | VoxelMorph, LessNet, Fourier-Net |
| SITReg | Symmetric, inverse-consistent, topology-preserving registration | SITReg |
| ConvexAdam, FireANTs, deedsBCV | Optimization-based and classical baselines | ConvexAdam, deedsBCV |
| TransMorph-LDDMM, LDDMMWrapper | Time-varying velocity fields (geodesic shooting) for any backbone | LDDMM |
| MultiMorph, TemplateCreation | Atlas construction | Atlas construction |
Conventions, transformation models, regularization, losses, metrics and a complete training recipe are
explained in the user guide.
Tutorials
Runnable training, inference and evaluation scripts live under tutorials/, one folder per
dataset and model (see the tutorial catalogue and
tutorials/README.md). Good starting points:
여기까지가 README 앞부분입니다. 나머지는 원문에서 읽을 수 있습니다. README 전문 보기 ↗
