Registration Playground

A registration network takes a moving and a fixed image and predicts a dense displacement field that warps the moving image onto the fixed one. Training balances image similarity against a smoothness (regularization) penalty on the displacement field. In most deep-learning registration models that penalty is spatially invariant: one weight for the whole image.

The spatially varying model, the formulation I introduced in Chen et al. (2025), instead learns a per-pixel weight λ(x) ∈ (0, 1] end to end. The penalty becomes w·λ(x)·|∇u(x)|², and a Beta prior with parameter α = 1 + κw keeps λ from collapsing to zero. Per pixel, the optimum is λ* = κ / |∇u|² (clipped to 1): where the deformation must be sharp and local, λ drops and the penalty relaxes; where tissue moves coherently, λ stays near 1. The prior strength slider sets κ, the regularization weight slider sets w. Because both models are hypernetwork-conditioned, a single network covers the whole range of hyperparameters, so every slider move is one forward pass rather than a retrained model.

The Jacobian determinant map shows local expansion (red) and compression (blue) of the deformation; pixels with |J| ≤ 0 are foldings, where the transformation is no longer invertible. Dice is computed over the 24 structures present in the slice labels.

Built with MIR

Everything here comes from my toolkit: the model, the losses, the training loop, and the export. The full recipe is a documented MIR tutorial, so you can train and deploy your own.

from MIR.models import HyperVFASPR, SpatialTransformer
from MIR.deformation_regularizer import logBeta, LocalGrad2d
import MIR.models.configs_VFA as cfgs

cfg = cfgs.get_VFA_default_config()
cfg.img_size, cfg.nb_hyp_units, cfg.initialize = (192, 160), 16, 1.0
model = HyperVFASPR(cfg, device='cuda')          # 3.8M parameters
warp  = SpatialTransformer((192, 160))

hyper = torch.rand(1, 2)                          # (kappa, w) codes in [0, 1]
flow, lam = model((moving, fixed), hyper)
loss = mse(warp(moving, flow), fixed) \
     + w * LocalGrad2d('l2')(flow, lam) \
     + logBeta()(lam, alpha=1 + kappa * w)

torch.onnx.export(model, ..., opset_version=16)   # then onnxruntime-web

• MIR on GitHub · documentation
• Paper: Unsupervised learning of spatially varying regularization for diffeomorphic image registration, Medical Image Analysis, 2025
• Hypernetwork conditioning follows HyperMorph (Hoopes et al., 2021)

Images: IXI dataset (CC BY-SA 3.0), T1-weighted, preprocessed and skull-stripped, test split only. The models are small (3.8M parameters) 2D networks trained for this demo with a mean-squared-error similarity; they are not the 3D models used in the papers. Inference uses ONNX Runtime Web (WebAssembly) and nothing leaves your device.