Merge pull request #2 from HiPCTProject/simple-test

Add a test for registration_rot

.pre-commit-config.yaml

@@ -17,6 +17,8 @@ repos:
     rev: v1.8.0
     hooks:
       - id: mypy
+        additional_dependencies:
+          - pytest
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v4.0.0-alpha.8
     hooks:

pyproject.toml

@@ -17,6 +17,12 @@ classifiers = [
     "Typing :: Typed",
 ]
 dependencies = [
+    "dask-image~=2023.08",
+    "itk~=5.3",
+    "numpy~=1.26",
+    "psutil~=5.9",
+    "scikit-image~=0.22",
+    "simpleitk~=2.3",
 ]
 description = "Code to register regions of interest with full organ datasets."
 dynamic = [
@@ -35,6 +41,7 @@ urls.homepage = "https://github.com/HiPCTProject/hipct-reg"
 test = [
     "pytest",
     "pytest-cov",
+    "tifffile==2024.1.30",
 ]
 dev = [
     "hipct-reg[test]",
@@ -58,6 +65,14 @@ paths.source = [
 
 [tool.mypy]
 explicit_package_bases = true
+disallow_untyped_calls = false
+
+[[tool.mypy.overrides]]
+module = [
+    "hipct_reg.tests"
+]
+strict = true
+warn_return_any = false
 
 [tool.pytest.ini_options]
 addopts = """
@@ -68,7 +83,7 @@ addopts = """
     --cov-report=xml
 """
 testpaths = [
-    "tests",
+    "src/hipct_reg/tests",
 ]
 
 [tool.ruff]

src/hipct_reg/ITK_registration.py

@@ -34,11 +34,11 @@ def registration_rot(
     fixed_image: sitk.Image,
     moving_image: sitk.Image,
     trans_point: npt.NDArray,
-    pt_fixed: npt.NDArray,
+    rotation_center_pix: npt.NDArray,
     zrot: float,
     angle_range: float,
     angle_step: float,
-    fiji: bool = True,
+    fiji: bool = False,
 ):
     """
     Parameters
@@ -47,8 +47,9 @@ def registration_rot(
         The images being registered.
     trans_point :
         Vector from [0, 0, 0] voxel in fixed image to [0, 0, 0] voxel in moving image.
-    pt_fixed :
-        Common point in the fixed image. In units of pixels.
+    rotation_center :
+        Point in the fixed image about which the moving image will be rotated.
+        In units of pixels.
     zrot :
         Initial rotation for the registration. In units of degrees.
     angle_range :
@@ -79,8 +80,8 @@ def registration_rot(
     R.SetOptimizerScalesFromPhysicalShift()
 
     offset = pixel_size_fixed * trans_point
-
-    rotation_center = pt_fixed * pixel_size_fixed
+    # Convert from pixels to physical size
+    rotation_center = rotation_center_pix * pixel_size_fixed
 
     theta_x = 0.0
     theta_y = 0.0
@@ -129,7 +130,7 @@ def command_iteration(
         rotation_center,
         moving_image,
         fixed_image,
-        pt_fixed,
+        rotation_center_pix,
     ):
         global metric
         metric.append(method.GetMetricValue())
@@ -153,7 +154,7 @@ def command_iteration(
             rotation_center,
             moving_image,
             fixed_image,
-            pt_fixed,
+            rotation_center_pix,
         ),
     )
 
@@ -309,7 +310,7 @@ def convertSI(fixed_image):
 
 
 def registration_sitk(
-    fixed_image, moving_image, trans_point, zrot, pt_fixed, fiji=True
+    fixed_image, moving_image, trans_point, zrot, pt_fixed, fiji=False
 ):
     pixel_size_fixed = fixed_image.GetSpacing()[0]
     pixel_size_moved = moving_image.GetSpacing()[0]
@@ -456,7 +457,7 @@ def command_iteration(method, pixel_size, trans_point):
 
 
 def registration_simpleElastix(
-    fixed_image, moving_image, trans_point, zrot, pt_fixed, fiji=True
+    fixed_image, moving_image, trans_point, zrot, pt_fixed, fiji=False
 ):
     print("\n-----------------")
     print("\nBeginning of initialization")

src/tests/test_registration.py

@@ -0,0 +1,112 @@
+"""An example set of tests."""
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+import SimpleITK as sitk
+import tifffile
+from skimage.data import binary_blobs
+from skimage.measure import block_reduce
+
+from hipct_reg.helpers import import_im
+from hipct_reg.ITK_registration import registration_rot
+
+PIXEL_SIZE_UM = 5
+BIN_FACTOR = 4
+ROI_OFFSET = 128
+ROI_SIZE = 64
+
+
+@pytest.fixture
+def rng() -> np.random.Generator:
+    """
+    Setup the default random number generator.
+    """
+    return np.random.default_rng(seed=1)
+
+
+def write_array_to_stack(array: npt.NDArray[Any], folder: Path) -> None:
+    """
+    Write a 3D array to a stack of images.
+    """
+    assert array.ndim == 3, "Input array must be 3D"
+    for i, plane in enumerate(array):
+        tifffile.imwrite(folder / f"{i}.tif", plane, photometric="minisblack")
+
+
+@pytest.fixture
+def ground_truth(rng: np.random.Generator) -> npt.NDArray[np.float32]:
+    """
+    Ground truth, high resolution data.
+    """
+    return binary_blobs(
+        length=512, n_dim=3, blob_size_fraction=0.01, volume_fraction=0.5, rng=rng
+    ).astype(np.float32)
+
+
+@pytest.fixture
+def full_organ_scan(
+    tmp_path: Path, ground_truth: npt.NDArray[np.float32]
+) -> sitk.Image:
+    """
+    Downsampled ground truth data, to mimic a full organ scan.
+    """
+    full_organ_scan = block_reduce(
+        ground_truth, (BIN_FACTOR, BIN_FACTOR, BIN_FACTOR), np.mean
+    )
+    full_organ_folder = tmp_path / "20.0um_full_organ"
+    full_organ_folder.mkdir()
+    write_array_to_stack(full_organ_scan, full_organ_folder)
+
+    return import_im(
+        str(full_organ_folder), pixel_size=PIXEL_SIZE_UM * BIN_FACTOR * 1000
+    )
+
+
+@pytest.fixture
+def roi_scan(tmp_path: Path, ground_truth: npt.NDArray[np.float32]) -> sitk.Image:
+    """
+    Sub-volume of ground truth data, to mimic ROI data.
+    """
+    roi_scan = ground_truth[
+        ROI_OFFSET : ROI_OFFSET + ROI_SIZE,
+        ROI_OFFSET : ROI_OFFSET + ROI_SIZE,
+        ROI_OFFSET : ROI_OFFSET + ROI_SIZE,
+    ]
+    roi_folder = tmp_path / "5.0um_roi"
+    roi_folder.mkdir()
+    write_array_to_stack(roi_scan, roi_folder)
+
+    return import_im(str(roi_folder), pixel_size=PIXEL_SIZE_UM * 1000)
+
+
+def test_registration_rot(full_organ_scan: sitk.Image, roi_scan: sitk.Image) -> None:
+    """
+    Test a really simple registration where the common point given is exactly the
+    correct point, and there is no rotation between the two datasets.
+    """
+    zrot = 0
+
+    trans_point = np.array([ROI_OFFSET, ROI_OFFSET, ROI_OFFSET]) / BIN_FACTOR
+    rotation_center = (
+        trans_point + np.array([ROI_SIZE, ROI_SIZE, ROI_SIZE]) / 2 / BIN_FACTOR
+    )
+
+    transform = registration_rot(
+        full_organ_scan,
+        roi_scan,
+        trans_point=trans_point,
+        rotation_center_pix=rotation_center,
+        zrot=zrot,
+        angle_range=360,
+        angle_step=2,
+    )
+
+    assert isinstance(transform, sitk.Euler3DTransform)
+    assert transform.GetAngleX() == 0
+    assert transform.GetAngleY() == 0
+    # This value should be close to zero
+    assert np.rad2deg(transform.GetAngleZ()) == pytest.approx(-2)