Add apply_transform pipeline for moving images and channels

README.md

@@ -25,7 +25,7 @@ Optional: If you want to hide all warnings during tests, append:
 
 This test generates deformed test data, runs the Snakemake workflow, and compares the final outputs against pre-computed reference results.
 
-The reference data will be downloaded automatically from this repo's release if available. If the download fails (e.g. the repository is private), manually download the reference data from [here](https://github.com/kreshuklab/matchmaker/releases/tag/test_data-v0.2) and place it under:
+The reference data will be downloaded automatically from this repo's release if available. If the download fails (e.g. the repository is private), manually download the reference data from [here](https://github.com/kreshuklab/matchmaker/releases/tag/test_data-v1.0) and place it under:
 
 ```
 examples/data/test_data/
@@ -80,8 +80,8 @@ mm.n5-utils.read_volume(...)
 ## Registration
 ### Input
 
-- **Fixed image**: 3D instance segmentation in `n5` + resolution 
-- **Moving image**: 3D instance segmentation `n5` + resolution 
+- **Fixed image**: 3D instance segmentation in `n5` + resolution
+- **Moving image**: 3D instance segmentation `n5` + resolution
 
 
 Expected image shape: ZYX
@@ -100,7 +100,7 @@ Expected image shape: ZYX
 **1. PCA pre-alignment**: alignment of fixed and moving image to the PCs \
  `prealignment.py --fixed_path ... --fixed_key ... --moving_path ... --moving_key ... --output_dir ... --mobie_export --dataset_name ...` \
 
-Outputs: 
+Outputs:
 - prealigned images: `{file_name}_prealigned.n5`
 - transformation matrixes
     - `{file_name}_fixed_T_prealignment.txt`
@@ -114,7 +114,7 @@ Outputs:
 **2. Rigid pre-alignment with Elastix** \
 `apply_rigid_elastix.py --fixed_path ... --fixed_key ... --moving_path ... --moving_key ... --output_dir ... --mobie_export --dataset_name ...` \
 
-Outputs: 
+Outputs:
 - rigid alinged moving image: `{file_name}_rigid_aligned.n5`
 - rigid transformation matrix (Elastix outputs): `result.0.mhd`, `result.0.raw`, `TransformParameters.0.txt`
 - logging file: `elastix_log_rigid.log`
@@ -143,4 +143,4 @@ with distance between keypoints in loss and rigidity penalty ---> **B-spline coe
 - **Moving image resampled to match fixed image** - only rigid - `n5`
 - **Moving image resampled to match fixed image** - deformable - `n5`
 
-Expected image shape: (C)ZYX
+Expected image shape: (C)ZYX

examples/deform_test_data.py

@@ -8,7 +8,7 @@
 
 from matchmaker.utils import (get_transformation_matrix, rotate_img, write_volume,
                                 plot_three_slices, plot_overlay, grid_sample3d, load_config,
-                                crop_to_bbox)
+                                crop_to_bbox, resample_volume)
 
 
 def remove_instances(seg, prob=0.05, seed=None):
@@ -32,33 +32,38 @@ def remove_instances(seg, prob=0.05, seed=None):
     return seg
 
 
-def save_volume(path, array, key="seg", chunks=(128, 512, 512), attributes={"resolution":[1,1,1]},
+def save_volume(path, array, key="seg", chunks=(128, 512, 512), resolution=[1,1,1],
                 save_tif=True):
     assert path.endswith(".n5")
 
-    write_volume(f=path, arr=array, key=key, chunks=chunks, attrs=attributes,)
+    write_volume(f=path, arr=array, key=key, chunks=chunks, attrs={"resolution":resolution,},)
 
     if save_tif:
         tif.imwrite(path.replace(".n5", ".tif"), array)
 
 
-def rigid_deform(fixed, angles):
+def rigid_deform(fixed, angles, voxel_spacing=None):
     if not isinstance(angles, (list, tuple)):
         raise TypeError()
 
     assert len(angles) == 3
+    iso_spacing = np.asarray([1, 1, 1], dtype=np.float32)
 
     center = np.array(fixed.shape) // 2
     rotation = tf3d.euler.euler2mat(
         *[np.deg2rad(angles[0]), np.deg2rad(angles[1]), np.deg2rad(angles[2])], axes="szyx"
     )
 
-    T, new_shape = get_transformation_matrix(fixed, center, rotation)
+    T, new_shape = get_transformation_matrix(fixed, center, rotation, iso_spacing)
     moving = rotate_img(fixed, T, output_shape=new_shape)
+
+    if not np.array_equal(voxel_spacing, iso_spacing):
+        moving = resample_volume(moving, iso_spacing, voxel_spacing)
+
     return moving
 
 
-def elastic_deform(volume, alpha=(1.,1.,1.), sigma=None, spacing=16, mode="nearest",
+def elastic_deform(volume, alpha=(1.,1.,1.), sigma=None, grid_spacing=16, mode="nearest",
                     align_corners=False, seed=None,):
     """
     Apply elastic deformation to a 3D volume.
@@ -67,7 +72,7 @@ def elastic_deform(volume, alpha=(1.,1.,1.), sigma=None, spacing=16, mode="neare
         volume (np.ndarray): Input volume of shape (D, H, W).
         alpha (tuple[float, float, float]): Displacement amplitude scaling factors.
         sigma (float): Gaussian smoothing std.
-        spacing (int | tuple[int, int, int]): Control point spacing (in voxels).
+        grid_spacing (int | tuple[int, int, int]): Control point spacing (in voxels).
         mode (str): Interpolation mode, "nearest" or "trilinear".
         align_corners (bool): Grid sampling alignment flag.
         seed (int | None): Random seed.
@@ -91,15 +96,15 @@ def elastic_deform(volume, alpha=(1.,1.,1.), sigma=None, spacing=16, mode="neare
     else:
         raise ValueError
 
-    if isinstance(spacing, int):
-        spacing = [spacing] * 3
-    elif isinstance(spacing, (list, tuple)):
-        assert len(spacing) == 3
+    if isinstance(grid_spacing, int):
+        grid_spacing = [grid_spacing] * 3
+    elif isinstance(grid_spacing, (list, tuple)):
+        assert len(grid_spacing) == 3
     else:
         raise ValueError
 
     if sigma is None:
-        sigma = (spacing[0] / 2, spacing[1] / 2, spacing[2] / 2)
+        sigma = (grid_spacing[0] / 2, grid_spacing[1] / 2, grid_spacing[2] / 2)
     elif isinstance(sigma, (int, float)):
         sigma = (sigma, sigma, sigma)
     elif isinstance(sigma, (list, tuple)):
@@ -108,13 +113,13 @@ def elastic_deform(volume, alpha=(1.,1.,1.), sigma=None, spacing=16, mode="neare
     else:
         raise ValueError
 
-    shape = (int(np.ceil(D/spacing[0])), int(np.ceil(H/spacing[1])), int(np.ceil(W/spacing[2])))
+    shape = (int(np.ceil(D/grid_spacing[0])), int(np.ceil(H/grid_spacing[1])), int(np.ceil(W/grid_spacing[2])))
     disp = np.random.randn(*shape, 3).astype(np.float32)
 
     disp = gaussian(disp, sigma=(*sigma, 0), mode="constant", preserve_range=True,)
     disp *= alpha
 
-    if (spacing[0] > 1) or (spacing[1] > 1) or (spacing[2] > 1):
+    if (grid_spacing[0] > 1) or (grid_spacing[1] > 1) or (grid_spacing[2] > 1):
         zoom_factors = (D / shape[0], H / shape[1], W / shape[2], 1.,)
         disp = zoom(disp, zoom_factors, order=1)
 
@@ -132,42 +137,50 @@ def normalize_axis(d):
     return sampled.astype(volume.dtype)
 
 
-def deform_test_data(cfg_path="", config=None, enable_elastic=False, alpha=0.9, sigma=2, spacing=16,
-                        rotate_angles_fixed=[20,345,30], rotate_angles_moving=[155,30,65],
-                        remove_p=0.05, seed=42, visualize=True):
+def deform_test_data(cfg_path="", config=None, enable_aniso=False, enable_elastic=False,
+                        alpha=0.9, sigma=2, grid_spacing=16, rotate_angles_fixed=[20,345,30],
+                        rotate_angles_moving=[155,30,65], remove_p=0.05, seed=42, visualize=True):
     if config is None:
         config = load_config(cfg_path)
 
     data_dir = Path(config["fixed_image"]["path"]).parent
     data_dir.mkdir(parents=True, exist_ok=True)
 
+    fixed_spacing = [config["fixed_image"]["z_res"], config["fixed_image"]["y_res"], config["fixed_image"]["x_res"]]
+    moving_spacing = [config["moving_image"]["z_res"], config["moving_image"]["y_res"], config["moving_image"]["x_res"]]
+    fixed_spacing = np.asarray(fixed_spacing, dtype=np.float32)
+    moving_spacing = np.asarray(moving_spacing, dtype=np.float32)
+
+    if enable_aniso:
+        assert not np.all(moving_spacing == 1)
+
     seg_fixed = tif.imread(config["fixed_image"]["source_path"])
+    seg_moving = seg_fixed.copy()
 
-    seg_fixed = rigid_deform(seg_fixed, angles=rotate_angles_fixed)
+    seg_fixed = rigid_deform(seg_fixed, rotate_angles_fixed, fixed_spacing)
     seg_fixed = crop_to_bbox(seg_fixed)
     print("Fixed volume shape", seg_fixed.shape)
 
-    seg_moving = seg_fixed.copy()
     if enable_elastic:
-        seg_moving = elastic_deform(seg_moving, alpha=alpha, sigma=sigma, spacing=spacing, seed=seed,)
+        seg_moving = elastic_deform(seg_moving, alpha=alpha, sigma=sigma, grid_spacing=grid_spacing, seed=seed,)
 
-    seg_moving = rigid_deform(seg_moving, angles=rotate_angles_moving)
+    seg_moving = rigid_deform(seg_moving, rotate_angles_moving, moving_spacing)
     seg_moving = remove_instances(seg_moving, prob=remove_p, seed=seed)
     seg_moving = crop_to_bbox(seg_moving)
     print("Moving volume shape", seg_moving.shape)
 
-    save_volume(config["fixed_image"]["path"].replace(".tif", ".n5"), seg_fixed)
-    save_volume(config["moving_image"]["path"].replace(".tif", ".n5"), seg_moving)
+    save_volume(config["fixed_image"]["path"].replace(".tif", ".n5"), seg_fixed, resolution=fixed_spacing.tolist())
+    save_volume(config["moving_image"]["path"].replace(".tif", ".n5"), seg_moving, resolution=moving_spacing.tolist())
 
     if visualize:
         plot_dir = data_dir / "plots"
         plot_dir.mkdir(parents=True, exist_ok=True)
-        plot_three_slices(seg_fixed, save_path=plot_dir/"seg_fixed.png")
+        iso_name = "aniso_" if enable_aniso else ""
+        elastic_name = "elastic" if enable_elastic else "rigid"
 
-        moving_name = "seg_elastic" if enable_elastic else "seg_rigid"
-        plot_three_slices(seg_moving, save_path=plot_dir/f"{moving_name}.png")
-        overlay_name = "elastic_overlay" if enable_elastic else "rigid_overlay"
-        plot_overlay(seg_fixed, seg_moving, save_path=plot_dir/f"{overlay_name}.png")
+        plot_three_slices(seg_fixed, save_path=plot_dir/f"seg_{iso_name}fixed.png")
+        plot_three_slices(seg_moving, save_path=plot_dir/f"seg_{iso_name}{elastic_name}.png")
+        plot_overlay(seg_fixed, seg_moving, save_path=plot_dir/f"{iso_name}{elastic_name}_overlay.png")
 
 
 if __name__ == "__main__":

examples/register_config_test_aniso_rigid.yaml

@@ -0,0 +1,36 @@
+fixed_image:
+  source_path: "examples/data/platy1_muscles_stardist_fixed.tif"
+  path: "examples/data/deformed_data/platy1_muscles_stardist_aniso_fixed.tif"
+  output_name: "fixed_image"
+  x_res: 1
+  y_res: 1
+  z_res: 1
+
+moving_image:
+  path: "examples/data/deformed_data/platy1_muscles_stardist_aniso_rigid.tif"
+  output_name: "moving_image"
+  x_res: 2
+  y_res: 1
+  z_res: 1
+  ref_path: "examples/data/test_data/moving_aniso_rigid_aligned.tif"
+  ref_url: "https://github.com/kreshuklab/matchmaker/releases/download/test_data-v0.3/moving_aniso_rigid_aligned.tif"
+
+log_dir: "data/test_rigid_registration"
+final_transform_path: "data/test_rigid_registration/final_transform.json"
+
+prealignment:
+  axis_orientation: "auto"
+
+coherent_point_drift:
+  w: 0.00001
+  lmd: 0.1
+  beta: 100
+  maxiter: 10 # Low for debugging purpose, should be 100-150
+
+ILP:
+  min_neighbours: 10
+  max_dist: 30
+
+mobie_export: True
+semantic_seg: True
+mobie_dataset_name: "platy1_muscles_stardist"

examples/register_config_test_rigid.yaml

@@ -7,17 +7,20 @@ fixed_image:
   z_res: 1
 
 moving_image:
-  path: "examples/data/deformed_data/platy1_muscles_stardist_moving.tif"
+  path: "examples/data/deformed_data/platy1_muscles_stardist_rigid.tif"
   output_name: "moving_image"
   x_res: 1
   y_res: 1
   z_res: 1
-  ref_path: "examples/data/test_data/moving_rigid_aligned.tif"
-  ref_url: "https://github.com/kreshuklab/matchmaker/releases/download/test_data-v0.2/moving_rigid_aligned.tif"
+  ref_path: "examples/data/test_data/moving_pointset_aligned.tif"
+  ref_url: "https://github.com/kreshuklab/matchmaker/releases/download/test_data-v1.0/moving_pointset_aligned.tif"
 
 log_dir: "data/test_rigid_registration"
 final_transform_path: "data/test_rigid_registration/final_transform.json"
 
+prealignment:
+  axis_orientation: "auto"
+
 coherent_point_drift:
   w: 0.00001
   lmd: 0.1

examples/register_config_test_rigid_apply_transform.yaml

@@ -0,0 +1,27 @@
+fixed_image:
+  input_path: "examples/data/deformed_data/platy1_muscles_stardist_fixed_rotated.tif"
+  input_key: "input"
+
+moving_images:
+  -
+    input_path: "data/test_rigid_registration/moving_image.n5"
+    input_key: "input"
+    output_path: "data/test_apply_transform/moving_image.n5"
+    output_key: "pointset_alignment_transform"
+    x_res: 1
+    y_res: 1
+    z_res: 1
+    interpolation_order: 0
+  -
+    input_path: "examples/data/deformed_data/platy1_muscles_stardist_rigid.tif"
+    input_key: "input"
+    output_path: "data/test_apply_transform/moving_image.tif"
+    output_key: "pointset_alignment_transform"
+    x_res: 1
+    y_res: 1
+    z_res: 1
+    interpolation_order: 0
+
+log_dir: "data/test_apply_transform"
+parameter_map_path: "data/test_rigid_registration/elastix_deformable_pointset_registration/TransformParameters.2.txt"
+prealignment_transform_path: "data/test_rigid_registration/svd_prealignment/svd_prealignment_transform.json"