torchreg is a tiny (~100 lines) PyTorch-based library for 2D and 3D image registration.
Affine Registration of two image tensors is done via:
from torchreg import AffineRegistration
# Load images as torch Tensors
small_alice = ... # Tensor with shape [1, 3 (color channel), 1024 (pixel), 1024 (pixel)]
big_alice = ... # Tensor with shape [1, 3 (color channel), 1024 (pixel), 1024 (pixel)]
# Intialize AffineRegistration
reg = AffineRegistration(is_3d=False)
# Run it!
moved_alice = reg(small_alice, big_alice)Multiresolution approach to save compute (per default 1/4 + 1/2 of original resolution for 500 + 100 iterations)
reg = AffineRegistration(scales=(4, 2), iterations=(500, 100))Choosing which operations (translation, rotation, zoom, shear) to optimize
reg = AffineRegistration(with_zoom=False, with_shear=False)Custom initial parameters
reg = AffineRegistration(zoom=torch.Tensor([[1.5, 2.]]))Custom dissimilarity functions and optimizers
def dice_loss(x1, x2):
dim = [2, 3, 4] if len(x2.shape) == 5 else [2, 3]
inter = torch.sum(x1 * x2, dim=dim)
union = torch.sum(x1 + x2, dim=dim)
return 1 - (2. * inter / union).mean()
reg = AffineRegistration(dissimilarity_function=dice_loss, optimizer=torch.optim.Adam)CUDA support (NVIDIA GPU)
moved_alice = reg(moving=big_alice.cuda(), static=small_alice.cuda())MPS support (Apple M1 or M2)
moved_alice = reg(moving=big_alice.to('mps'), static=small_alice.to('mps'))After the registration is run, you can apply it to new images (coregistration)
another_moved_alice = reg.transform(another_alice, shape=(256, 256))with desired output shape.
You can access the affine
affine = reg.get_affine()and the four parameters (translation, rotation, zoom, shear)
translation = reg.parameters[0]
rotation = reg.parameters[1]
zoom = reg.parameters[2]
shear = reg.parameters[3]pip install torchregThere are three example notebooks:
- examples/basics.ipynb shows the basics by using small cubes/squares as image data
- examples/images.ipynb shows how to register alice_big.jpg to alice_small.jpg
- examples/mri.ipynb shows how to register MR images (Nifti files) including co-, parallel and multimodal registration
If you want to know how the core of this package works, read the blog post!