Symmetric diffeomorphic and affine image registration methods in PyTorch and JAX
Project description
syntx
syntx is a high-performance Python package focusing on symmetric diffeomorphic (SyN) and affine image registration methods, built on top of PyTorch and JAX for GPU/MPS acceleration and auto-differentiation capabilities.
Ported from the registration modules of the sulceye package, syntx is designed for distribution on PyPI and works seamlessly with standard medical image types, particularly ANTsImage from the antspyx library.
⚠️ Disclaimer & Differences from ants.registration
[!IMPORTANT] Validation Status: The deep-learning feature-space similarity metrics (VGG19, DINOv2, Swin UNETR) in this repository are experimental and have not been deeply validated on large-scale clinical cohorts. They are intended strictly for research and exploration.
Key Differences from
ants.registration:
- GPU Acceleration: Unlike standard
ants.registration(which runs on CPU via ITK C++), Syntx supports PyTorch and JAX optimization backends for fast GPU/MPS execution.- Optimizers: Syntx uses Adam/Rprop for the affine stage and greedy composition steps scaling by ITK-style CFL (Courant-Friedrichs-Lewy) max voxel displacement bounds, while ANTs relies on C++ variants of L-BFGS or regularized gradient descent.
- Velocity-Field & Elastic Smoothing: Separable Gaussian filters are implemented natively in JAX/PyTorch to perform fluid-like smoothing of update fields and elastic-like smoothing of composed fields, matching ITK's Gaussian regularization on the GPU.
- Multi-Resolution Pyramid: Downsampling is performed dynamically using bilinear/trilinear grid interpolation in PyTorch/JAX to build image pyramids, rather than ITK's C++ downsampling filters.
- Feature-Space Metrics: Similarity is evaluated on multi-scale feature representations (from vision transformers or CNNs) via zero-copy DLPack autograd sharing, rather than raw intensity maps.
Key Features
- Auto-Differentiation Backends: Choose between
'pytorch'and'jax'for core computations. - Symmetric Normalization (SyN): Fully symmetric greedy optimization matching classic ITK/ANTs SyN implementations.
- Interoperability: Seamless conversions between PyTorch/JAX coordinate spaces and ITK physical coordinate matrices.
- Direct PyPy/PyPI Packaging: Implemented cleanly with minimum external dependencies.
Installation
To install syntx locally from the repository:
pip install -e .
Dependencies
numpyscipymatplotlibantspyxtorchjaxjaxlib
Usage Example
syntx exposes a high-level syn and registration API that mirrors ants.registration:
import ants
import syntx
# Load ANTs images
fixed = ants.image_read( ants.get_data('r16') )
moving = ants.image_read( ants.get_data('r64') )
# Run registration using PyTorch (default)
result = syntx.syn(
fixed=fixed,
moving=moving,
type_of_transform='SyNTo',
backend='pytorch',
reg_iterations=[100, 100, 50],
affine_iterations=[100, 50, 20],
)
# Access the warped moving output image
warped_moving = result['warpedmovout']
# Access transform files (saved to temporary paths for ANTs compatibility)
forward_transforms = result['fwdtransforms']
inverse_transforms = result['invtransforms']
For JAX backend acceleration:
result = syntx.syn(
fixed=fixed,
moving=moving,
type_of_transform='SyNTo',
backend='jax',
reg_iterations=[100, 100, 50],
affine_iterations=[100, 50, 20],
)
Running the Examples and Generating Reports
An example comparing classic ANTs, PyTorch, and JAX registration is included under examples/. It generates a comparison report summarizing Mutual Information, Jacobian Determinants (topological safety), and Execution Speed.
To run the comparison:
python examples/generate_ants_2d_comparison_report.py
This generates an HTML report under reports/ants_2d_syn_comparison.html.
Running Tests
Tests can be executed via pytest:
pytest
Makefile Automation
A Makefile is included to automate standard development tasks:
- Install (install package in editable mode):
make install - Test (run test suite in Fast mode, skipping slow 3D registrations, and printing a code coverage table):
make test
- Test All (run the full test suite including slow 3D registrations, with coverage):
make test-all - Clean (remove build artifacts, cached directories, and temporary files):
make clean - Release (clean, build sdist and wheel packages, and upload to PyPI using twine):
make release
It automatically detects and prioritizes the active python virtual environment (VIRTUAL_ENV).
Release
make clean
python -m build .
python -m twine upload --config-file ~/.pypirc dist/*
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file syntx-0.1.16.tar.gz.
File metadata
- Download URL: syntx-0.1.16.tar.gz
- Upload date:
- Size: 92.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.12.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
b660aa33db44d9d8bcab2ecf4f4de3ad7b27ea83b384b6a87bb699a1b6e205fa
|
|
| MD5 |
53d9b11193672f15c1def1af303711b5
|
|
| BLAKE2b-256 |
b7076b9ec1e5e9c84d506c3db2a8792b09a23317456ac98195765e5c026a9de0
|
File details
Details for the file syntx-0.1.16-py3-none-any.whl.
File metadata
- Download URL: syntx-0.1.16-py3-none-any.whl
- Upload date:
- Size: 68.5 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.12.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a2987423c7e814c27608cfc9e548495f98e4232cbcb803ac035fb148cd1a93e0
|
|
| MD5 |
c91ffe4bb4459f09e93907ffbbb44fa6
|
|
| BLAKE2b-256 |
b5ab1f48e09e96e3bd7be6d2b7e71ce1f115ec86099ec06dc7841b13d6c4c992
|