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