GuidedMorph: Two-Stage Deformable Registration for Breast MRI

Official implementation of GuidedMorph: Two-Stage Deformable Registration for Breast MRI by Yaqian Chen, Hanxue Gu, Haoyu Dong, Qihang Li, Yuwen Chen, Nicholas Konz, Lin Li, and Maciej Mazurowski.

Important model compatibility note

The released Guided-Trans-DT checkpoint was trained with a two-stage TransMorph architecture. Earlier versions of this repository accidentally provided a VoxelMorph (VxmDense_2) model in the default training and inference scripts. Those definitions cannot load the released weights.

This version contains the matching TransMorph implementation and always loads weights with strict=True. The recovered published checkpoint was verified with:

• 333 state-dict tensors;
• 169,921,974 stored tensor values;
• zero missing keys, unexpected keys, or shape mismatches.

Run the included checker before inference:

python scripts/validate_checkpoint.py weights/final.pth.tar

Successful validation ends with:

Compatibility: OK (strict load; all keys and shapes match)

Installation

The published environment uses Python 3.9, PyTorch 2.1.2, CUDA 11.8, and the versions recorded in environment.yml.

git clone https://github.com/mazurowski-lab/GuidedMorph.git
cd GuidedMorph
conda env create -f environment.yml
conda activate guidedmorph

For an existing Python 3.9 environment, install the pinned Python packages with:

pip install -r requirements.txt

A CUDA GPU is strongly recommended. The model accepts volumes of shape 128 × 256 × 256 and is large enough that CPU inference is slow.

Download weights and example data

Download the Guided-Trans-DT final.pth.tar checkpoint and the ISPY2 external test set from the project Google Drive folder. Place the checkpoint at:

weights/final.pth.tar

Weights, experiment outputs, and result files are intentionally excluded from Git because of their size.

Quick smoke-test demo

After placing the checkpoint at weights/final.pth.tar, generate a deterministic synthetic case and run one end-to-end inference:

python scripts/create_synthetic_case.py
python infer.py \
  --checkpoint weights/final.pth.tar \
  --test-dir demo/test \
  --output-dir results/synthetic \
  --device cuda:0

The first command creates demo/test/synthetic.pkl. The second command strictly loads the released checkpoint, extracts the guidance maps, runs both TransMorph stages, evaluates the synthetic mask, and writes results/synthetic/metrics.json. This generated case is intended only to verify the installation and inference pipeline; it is not a benchmark or a substitute for medical data.

Data format

Each case is a trusted Python pickle containing a tuple or list. Pickle files can execute code while loading, so do not use files from untrusted sources.

Training cases contain:

(moving_image, fixed_image)

Inference cases may also contain any number of paired binary masks:

(
    moving_image,
    fixed_image,
    moving_mask_1,
    fixed_mask_1,
    moving_mask_2,
    fixed_mask_2,
    # ...
)

Every image and mask must have spatial shape (128, 256, 256). Arrays may be stored as [D, H, W], [1, D, H, W], or [1, 1, D, H, W]. Images are converted to float32 by the loaders.

Suggested layout:

demo/
├── train/
│   ├── case001.pkl
│   └── case002.pkl
└── test/
    ├── case101.pkl
    └── case102.pkl

Inference

Edit infer_config.json or override its most common settings on the command line:

python infer.py \
  --checkpoint weights/final.pth.tar \
  --test-dir demo/test \
  --device cuda:0

The default configuration is:

{
  "test_dir": "demo/test",
  "checkpoint": "weights/final.pth.tar",
  "output_dir": "results",
  "device": "auto",
  "feature_backend": "skimage",
  "mask_warp": "distance",
  "save_outputs": false
}

feature_backend: "skimage" reproduces the feature-extraction path used for the published training run. "torch" is a faster GPU implementation, but its output is not numerically identical to the original scikit-image pipeline.

The command writes per-case timing and Dice statistics to results/metrics.json. Pass --save-outputs to additionally save the warped image, warped guidance map, and displacement field for every case as compressed .npz files.

Useful overrides:

# Automatically use CUDA when available, otherwise CPU
python infer.py --device auto

# Faster experimental guidance extraction
python infer.py --feature-backend torch

# Save warped volumes and flows
python infer.py --save-outputs

Training

Training now uses the same two-stage TransMorph architecture as the released checkpoint:

python train.py --config config.json

The default configuration is:

{
  "train_dir": "demo/train",
  "output_dir": "experiments/guided_trans_dt",
  "device": "auto",
  "batch_size": 1,
  "num_workers": 4,
  "epochs": 500,
  "save_every": 2,
  "learning_rate": 0.0005,
  "loss_weights": [1.0, 1.0, 0.08],
  "similarity_loss": "mse",
  "feature_backend": "skimage",
  "resume": null,
  "seed": 42
}

Resume training with either the config or a CLI override:

python train.py --resume experiments/guided_trans_dt/final.pth.tar

Checkpoints are written atomically to experiments/guided_trans_dt/final.pth.tar. Each checkpoint contains the epoch, model state, optimizer state, and training configuration.

Repository structure

GuidedMorph/
├── models/
│   ├── TransMorph.py          # checkpoint-compatible two-stage model
│   ├── configs_TransMorph.py  # published 128×256×256 configuration
│   └── voxelmorph.py          # retained legacy baseline
├── data/                      # pickle datasets and transformations
├── checkpoint.py             # strict checkpoint loading and diagnostics
├── feature_extract.py        # original and PyTorch guidance extraction
├── infer.py                   # TransMorph inference CLI
├── train.py                   # TransMorph training CLI
├── train_vxm.py               # legacy VoxelMorph baseline
├── scripts/
│   ├── create_synthetic_case.py
│   └── validate_checkpoint.py
├── config.json
├── infer_config.json
├── environment.yml
└── requirements.txt

Troubleshooting

Missing or unexpected unet_model keys

You are loading the Guided-Trans-DT weights into the old VoxelMorph model. Use models.build_guidedmorph() or the provided infer.py script.

Input shape error

The published checkpoint is tied to (128, 256, 256). Apply the same preprocessing used by GuidedMorph before inference; silently resizing inside the model would change the registration problem and is therefore not performed.

Out-of-memory error

Use batch size 1, close other GPU jobs, and avoid --save-outputs unless needed. The training checkpoint also contains Adam optimizer state, so loading it requires more host memory than the model weights alone.

Model attribution

The Transformer backbone is adapted from TransMorph, and the spatial transformer follows VoxelMorph. Source attribution is also preserved in models/TransMorph.py.

Citation

If you use this code or model, please cite:

@article{chen2025guidedmorph,
  title={GuidedMorph: Two-Stage Deformable Registration for Breast MRI},
  author={Chen, Yaqian and Gu, Hanxue and Dong, Haoyu and Li, Qihang and Chen, Yuwen and Konz, Nicholas and Li, Lin and Mazurowski, Maciej},
  journal={arXiv preprint arXiv:2505.13414},
  year={2025}
}

For questions, contact Yaqian Chen or open a GitHub issue.