Initial commit

Author: codingfisch

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2023 Lukas Fisch
+Copyright (c) 2023 codingfisch
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

README.md

@@ -1,2 +1,91 @@
 # torchreg
-Lightweight image registration library using PyTorch
+
+torchreg is a tiny (~100 lines) PyTorch-based library for 2D and 3D image registration.
+
+<p float="left", align="center">
+<img src="https://github.com/codingfisch/torchreg/blob/main/examples/alice_big.jpg" width="256"/>
+<img src="https://github.com/codingfisch/torchreg/blob/main/examples/alice_small.jpg" width="256"/>
+<img src="https://github.com/codingfisch/torchreg/assets/55840648/dbf414cc-75e5-477c-9794-32f97a16ea21" width="256"/>
+</p>
+
+## Usage
+Affine Registration of two image tensors is done via:
+```python
+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)
+```
+
+## Features
+
+Multiresolution approach to save compute (per default 1/4 + 1/2 of original resolution for 500 + 100 iterations)
+```python
+reg = AffineRegistration(scales=(4, 2), iterations=(500, 100))
+```
+Choosing which operations (translation, rotation, zoom, shear) to optimize
+```python
+reg = AffineRegistration(with_zoom=False, with_shear=False)
+```
+Custom initial parameters
+```python
+reg = AffineRegistration(zoom=torch.Tensor([[1.5, 2.]]))
+```
+Custom dissimilarity functions and optimizers
+```python
+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(dissimilairity_function=dice_loss, optimizer=torch.optim.Adam)
+```
+CUDA support (NVIDIA GPU)
+```python
+moved_alice = reg(moving=big_alice.cuda(), static=small_alice.cuda())
+```
+MPS support (Apple M1 or M2)
+```python
+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)
+```python
+another_moved_alice = reg.transform(another_alice, shape=(256, 256))
+```
+with desired output shape.
+
+You can access the affine
+```python
+affine = reg.get_affine()
+```
+and the four parameters (translation, rotation, zoom, shear)
+```python
+translation = reg.parameters[0]
+rotation = reg.parameters[1]
+zoom = reg.parameters[2]
+shear = reg.parameters[3]
+```
+
+## Installation
+```bash
+pip install torchreg
+```
+
+## Examples/Tutorials
+
+There are three example notebooks:
+
+- [examples/basics.ipynb](https://github.com/codingfisch/torchreg/blob/main/examples/basic.ipynb) shows the basics by using small cubes/squares as image data
+- [examples/images.ipynb](https://github.com/codingfisch/torchreg/blob/main/examples/image.ipynb) shows how to register alice_big.jpg to alice_small.jpg
+- [examples/mri.ipynb](https://github.com/codingfisch/torchreg/blob/main/examples/mri.ipynb) shows how to register MR images (Nifti files) including co-, parallel and multimodal registration
+
+## Background
+
+If you want to know how the core of this package works, read [the blog post]()!

examples/alice_big.jpg

@@ -0,0 +1,28 @@
+[tool.poetry]
+name = 'torchreg'
+version = '0.0.1'
+description = 'Lightweight image registration library using PyTorch'
+authors = ['codingfisch <[email protected]>']
+license = 'MIT'
+readme = 'README.md'
+repository = 'https://github.com/codingfisch/torchreg'
+classifiers = [
+    'Programming Language :: Python :: 3',
+    'Operating System :: OS Independent',
+    'Intended Audience :: Science/Research'
+]
+
+
+[tool.poetry.dependencies]
+python = '^3.0'
+torch = '*'
+tqdm = '*'
+
+
+[tool.poetry.group.test.dependencies]
+pytest = '*'
+
+
+[build-system]
+requires = ['poetry-core']
+build-backend = 'poetry.core.masonry.api'

examples/alice_small.jpg

@@ -0,0 +1,61 @@
+import torch
+from unittest import TestCase
+from torchreg import AffineRegistration
+from torchreg.affine import compose_affine, affine_transform, init_parameters, _check_parameter_shapes
+
+
+class TestAffineRegistration(TestCase):
+    def test_fit(self):
+        for batch_size in [1, 2]:
+            for n_dim in [2, 3]:
+                reg = AffineRegistration(scales=(1,), is_3d=n_dim == 3, learning_rate=1e-1, verbose=False)
+                moving = synthetic_image(batch_size, n_dim, shift=1)
+                static = synthetic_image(batch_size, n_dim, shift=0)
+                fitted_moved = reg(moving, static, return_moved=True)
+                fitted_affine = reg.get_affine()
+                affine = torch.stack(batch_size * [torch.eye(n_dim + 1)[:n_dim]])
+                affine[:, -1, -1] += -1/3
+                self.assertTrue(torch.allclose(fitted_affine, affine, atol=1e-2))
+                moved = affine_transform(moving, affine)
+                self.assertTrue(torch.allclose(fitted_moved, moved, atol=1e-2))
+
+    def test_affine_transform(self):
+        for batch_size in [1, 2]:
+            for n_dim in [2, 3]:
+                moving = synthetic_image(batch_size, n_dim, shift=1)
+                static = synthetic_image(batch_size, n_dim, shift=0)
+                affine = torch.stack(batch_size * [torch.eye(n_dim + 1)[:n_dim]])
+                affine[:, -1, -1] += -1/3
+                moved = affine_transform(moving, affine)
+                self.assertTrue(torch.allclose(moved, static, atol=1e-6))
+
+    def test_init_parameters(self):
+        for batch_size in [1, 2]:
+            for is_3d in [False, True]:
+                params = init_parameters(is_3d=is_3d, batch_size=batch_size)
+                self.assertIsInstance(params, list)
+                self.assertEqual(len(params), 4)
+                for param in params:
+                    self.assertTrue(isinstance(param, torch.nn.Parameter))
+                _check_parameter_shapes(*params, is_3d=is_3d, batch_size=batch_size)
+
+    def test_compose_affine(self):
+        for batch_size in [1, 2]:
+            for n_dim in [2, 3]:
+                translation = torch.zeros(batch_size, n_dim)
+                rotation = torch.stack(batch_size * [torch.eye(n_dim)])
+                zoom = torch.ones(batch_size, n_dim)
+                shear = torch.zeros(batch_size, n_dim)
+                affine = compose_affine(translation, rotation, zoom, shear)
+                id_affine = torch.stack(batch_size * [torch.eye(n_dim + 1)[:n_dim]])
+                self.assertTrue(torch.equal(affine, id_affine))
+
+
+def synthetic_image(batch_size, n_dim, shift):
+    shape = [batch_size, 1, 7, 7, 7][:2 + n_dim]
+    x = torch.zeros(*shape)
+    if n_dim == 3:
+        x[:, :, 2 - shift:5 - shift, 2:5, 2:5] = 1
+    else:
+        x[:, :, 2 - shift:5 - shift, 2:5] = 1
+    return x

examples/basic.ipynb

@@ -0,0 +1 @@
+from .affine import AffineRegistration

examples/image.ipynb

@@ -0,0 +1,103 @@
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+
+class AffineRegistration:
+    def __init__(self, scales=(4, 2), iterations=(500, 100), is_3d=True, learning_rate=1e-2,
+                 verbose=True, dissimilarity_function=torch.nn.MSELoss(), optimizer=torch.optim.Adam,
+                 init_translation=None, init_rotation=None, init_zoom=None, init_shear=None,
+                 with_translation=True, with_rotation=True, with_zoom=True, with_shear=False,
+                 align_corners=True, interp_mode=None, padding_mode='border'):
+        self.scales = scales
+        self.iterations = iterations[:len(scales)]
+        self.is_3d = is_3d
+        self.learning_rate = learning_rate
+        self.verbose = verbose
+        self.dissimilarity_function = dissimilarity_function
+        self.optimizer = optimizer
+        self.inits = (init_translation, init_rotation, init_zoom, init_shear)
+        self.withs = (with_translation, with_rotation, with_zoom, with_shear)
+        self.align_corners = align_corners
+        self.interp_mode = 'trilinear' if is_3d else 'bilinear' if interp_mode is None else interp_mode
+        self.padding_mode = padding_mode
+        self._parameters = None
+
+    def __call__(self, moving, static, return_moved=True):
+        if len(moving.shape) - 4 != self.is_3d or len(static.shape) - 4 != self.is_3d:
+            raise ValueError(f'Expected moving and static to be {4 + self.is_3d}D Tensors (2 + Spatial Dims.). '
+                             f'Got size {moving.shape} and {static.shape}.')
+        if moving.shape != static.shape:
+            raise ValueError(f'Expected moving and static to have the same size. '
+                             f'Got size {moving.shape} and {static.shape}.')
+
+        self._parameters = init_parameters(self.is_3d, len(static), static.device, *self.withs, *self.inits)
+        interp_kwargs = {'mode': self.interp_mode, 'align_corners': self.align_corners}
+        moving_ = F.interpolate(moving, static.shape[2:], **interp_kwargs)
+        for scale, iters in zip(self.scales, self.iterations):
+            moving_small = F.interpolate(moving_, scale_factor=1 / scale, **interp_kwargs)
+            static_small = F.interpolate(static, scale_factor=1 / scale, **interp_kwargs)
+            self._fit(moving_small, static_small, iters)
+        return self.transform(moving, static.shape[2:]).detach() if return_moved else None
+
+    def _fit(self, moving, static, iterations):
+        optimizer = self.optimizer(self._parameters, self.learning_rate)
+        progress_bar = tqdm(range(iterations), disable=not self.verbose)
+        for self.iter in progress_bar:
+            optimizer.zero_grad()
+            moved = self.transform(moving, static.shape[2:], with_grad=True)
+            loss = self.dissimilarity_function(moved, static)
+            progress_bar.set_description(f'Shape: {[*static.shape]}; Dissimiliarity: {loss.item()}')
+            loss.backward()
+            optimizer.step()
+
+    def transform(self, moving, shape=None, with_grad=False):
+        affine = self.get_affine(with_grad)
+        return affine_transform(moving, affine, shape, self.interp_mode, self.padding_mode, self.align_corners)
+
+    def get_affine(self, with_grad=False):
+        affine = compose_affine(*self._parameters)
+        return affine if with_grad else affine.detach()
+
+
+def affine_transform(x, affine, shape=None, mode='bilinear', padding_mode='border', align_corners=True):
+    shape = x.shape[2:] if shape is None else shape
+    grid = F.affine_grid(affine, [len(x), len(shape), *shape], align_corners)
+    sample_mode = 'bilinear' if mode == 'trilinear' else mode  # grid_sample converts 'bi-' to 'trilinear' internally
+    return F.grid_sample(x, grid, sample_mode, padding_mode, align_corners)
+
+
+def init_parameters(is_3d=True, batch_size=1, device='cpu', with_translation=True, with_rotation=True, with_zoom=True,
+                    with_shear=True, init_translation=None, init_rotation=None, init_zoom=None, init_shear=None):
+    _check_parameter_shapes(init_translation, init_rotation, init_zoom, init_shear, is_3d, batch_size)
+    n_dim = 2 + is_3d
+    translation = torch.zeros(batch_size, n_dim).to(device) if init_translation is None else init_translation
+    rotation = torch.stack(batch_size * [torch.eye(n_dim)]).to(device) if init_rotation is None else init_rotation
+    zoom = torch.ones(batch_size, n_dim).to(device) if init_zoom is None else init_zoom
+    shear = torch.zeros(batch_size, n_dim).to(device) if init_shear is None else init_shear
+    params = [translation, rotation, zoom, shear]
+    with_grad = [with_translation, with_rotation, with_zoom, with_shear]
+    return [torch.nn.Parameter(param, requires_grad=grad) for param, grad in zip(params, with_grad)]
+
+
+def compose_affine(translation, rotation, zoom, shear):
+    _check_parameter_shapes(translation, rotation, zoom, shear, zoom.shape[-1] == 3, zoom.shape[0])
+    square_matrix = torch.diag_embed(zoom)
+    if zoom.shape[-1] == 3:
+        square_matrix[..., 0, 1:] = shear[..., :2]
+        square_matrix[..., 1, 2] = shear[..., 2]
+    else:
+        square_matrix[..., 0, 1] = shear[..., 0]
+    square_matrix = rotation @ square_matrix
+    return torch.cat([square_matrix, translation[:, :, None]], dim=-1)
+
+
+def _check_parameter_shapes(translation, rotation, zoom, shear, is_3d=True, batch_size=1):
+    n_dim = 2 + is_3d
+    params = {'translation': translation, 'rotation': rotation, 'zoom': zoom, 'shear': shear}
+    for name, param in params.items():
+        if param is not None:
+            desired_shape = (batch_size, n_dim, n_dim) if name == 'rotation' else (batch_size, n_dim)
+            if param.shape != desired_shape:
+                raise ValueError(f'Expected {name} to be size {desired_shape} since batch_size is {batch_size} '
+                                 f'and is_3d is {is_3d} -> {2 + is_3d} dimensions. Got size {param.shape}.')

examples/mri.ipynb

@@ -0,0 +1,8 @@
+import torch
+
+
+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()