Dev

.gitattributes

@@ -0,0 +1,5 @@
+.n5 filter=lfs diff=lfs merge=lfs -text
+.tif filter=lfs diff=lfs merge=lfs -text
+.n5/** filter=lfs diff=lfs merge=lfs -text
+examples/data filter=lfs diff=lfs merge=lfs -text
+examples/data/** filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -172,3 +172,14 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# test results
+examples/data/test/
+
+# huge data
+*.tif
+*.n5/
+*.png
+*.ome.zarr/
+
+.DS_Store

README.md

@@ -1,4 +1,4 @@
-# matchmaker
+# 💞 Matchmaker
 Tool for segmentation-based deformable registration and object matching
 
 
@@ -7,10 +7,8 @@ Tool for segmentation-based deformable registration and object matching
 ### Input
 
 - **Fixed image**: 3D instance segmentation in `n5` + resolution 
-
 - **Moving image**: 3D instance segmentation `n5` + resolution 
 
-- **Registration parameters yaml**: parameters of the `matchmaker` pipeline 
 
 Expected image shape: ZYX
 
@@ -19,17 +17,46 @@ Expected image shape: ZYX
 - **QC plots**
 - **Files with all transforms**
 - **Table of correspondence between instances in moving and fixed instance segmentations**
+- **Logging file: `registration.log`**
+- **Optional: Mobie project saved at `{output_dir}/mobie_project/`**
 
 
 ### Registration steps
 
-0. Create point cloud
-1. PCA pre-alignment ---> **rigid transform matrix**
-2. Manual input: should the image be flipped?
-3. Rigid pre-alignment with Elastix OR with rigid CPD ---> **rigid transform matrix**
-4. Coherent point drift
-5. Matching points with mixed integer programming
-6. Deformable registration with Elastix with distance between keypoints in loss and rigidity penalty ---> **B-spline coefficients**
+**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: 
+- prealigned images: `{file_name}_prealigned.n5`
+- transformation matrixes
+    - `{file_name}_fixed_T_prealignment.txt`
+    - `{file_name}_moving_T_prealignment.txt` (maybe final one is `moving_T_prealignment.txt`, couldn't figure this out)
+- plots:
+    - slice per dimension before pre-alignment: `{file_name}_fixed.png`, `{file_name}_moving.png`
+    - slice per dimension after pre-alignment: `{file_name}_fixed_prealigned.png`, `{file_name}_moving_prealigned.png`
+    - overlay of slice per dimension after pre-alignment: `overlay_prealignment.png`
+    - intensity profiles per axis and volume: `fixed_intensity_profile_{axis}.png`, `moving_intensity_profile_{axis}.png`
+
+**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: 
+- 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`
+- plots:
+    - `intersample_segm_overlay_before_alignment.png`
+    - `intersample_segm_rigid_alignment_semantic.png`
+
+
+**3. Coherent point drift**
+
+**4. Matching points with mixed integer programming**
+
+**5. Deformable registration with Elastix** \
+with distance between keypoints in loss and rigidity penalty ---> **B-spline coefficients**
+
+
 
 
 ## Apply registration to other images

environment.yml

@@ -16,6 +16,7 @@ dependencies:
     - optuna
     - proxsuite
     - cvxpy
+    - transforms3d
 
     - ipykernel
     - seaborn
@@ -29,5 +30,6 @@ dependencies:
         - pyscipopt
         - click
         - pyyaml
+        - itk
 
 

examples/deform_test_data.py

@@ -0,0 +1,90 @@
+import numpy as np
+import tifffile as tif
+import transforms3d as tf3d
+
+from matchmaker.transform_utils import (
+    downscale_seg,
+    get_transformation_matrix,
+    rotate_img,
+    crop_to_bbox,
+)
+from matchmaker.vis import plot_three_slices, plot_overlay
+from matchmaker.n5_utils import write_volume
+
+
+def remove_instances(seg, prob=0.05):
+    # Get all unique instance IDs, excluding background (assumed to be 0)
+    instance_ids = np.unique(seg)
+    instance_ids = instance_ids[instance_ids != 0]
+
+    # Randomly select 10% of the instance IDs
+    num_to_remove = int(len(instance_ids) * prob)
+    print(f"Number of instances to remove: {num_to_remove}")
+    selected_ids = np.random.choice(instance_ids, size=num_to_remove, replace=False)
+
+    # Create a mask for the selected IDs and set them to 0
+    mask = np.isin(seg, selected_ids)
+    seg[mask] = 0
+
+    print(f"Number of instances left: {len(np.unique(seg))}")
+    return seg
+
+
+def main():
+    seg = tif.imread("data/platy1_muscles_stardist.tif")
+
+    # downsample image
+    factor = 4
+    seg = downscale_seg(seg, factor)
+    seg_fixed = crop_to_bbox(seg)
+    print("Cropped shape", seg_fixed.shape)
+
+    # save downsampled, fixed image
+    attributes = {"resolution": [1, 1, 1]}
+    write_volume(
+        f="./data/platy1_muscles_stardist_fixed.n5",
+        arr=seg_fixed,
+        key="seg",
+        chunks=(128, 512, 512),
+        attrs=attributes,
+    )
+
+    # save also as tiff
+    tif.imwrite("./data/platy1_muscles_stardist_fixed.tif", seg_fixed)
+
+    # rotate image
+    center = np.array(seg_fixed.shape) // 2
+    rotation = tf3d.euler.euler2mat(
+        *[np.deg2rad(155), np.deg2rad(30), np.deg2rad(65)], axes="szyx"
+    )
+
+    T, new_shape = get_transformation_matrix(
+        seg_fixed, center, rotation, save_path="./data/transformation_matrix.txt"
+    )
+    seg_moving = rotate_img(seg_fixed, T, output_shape=new_shape)
+
+    # randomly remove instances
+    probability = 0.05
+    seg_moving = remove_instances(seg_moving, prob=probability)
+
+    # save moving image
+    attributes = {"resolution": [1, 1, 1]}
+    write_volume(
+        f="./data/platy1_muscles_stardist_moving.n5",
+        arr=seg_moving,
+        key="seg",
+        chunks=(128, 512, 512),
+        attrs=attributes,
+    )
+
+    # save also as tiff
+    tif.imwrite("./data/platy1_muscles_stardist_moving.tif", seg_moving)
+
+    # visualize
+    plot_three_slices(seg_moving, save_path="./data/plots/seg_moving.png")
+    plot_three_slices(seg_fixed, save_path="./data/plots/seg_fixed.png")
+    plot_overlay(seg_fixed, seg_moving, save_path="./data/plots/seg_overlay.png")
+
+
+if __name__ == "__main__":
+    main()

examples/mobie_dataset_clean_up.py

@@ -0,0 +1,35 @@
+import json
+
+# clean up dataset.json
+input_path = "./data/test/mobie_project/platy1_muscles_stardist/dataset.json"
+
+# Define the key you want to delete
+key_to_delete = "platy1_muscles_stardist_moving_prealigned_rigid_aligned"
+
+# Load the JSON file
+with open(input_path, "r") as file:
+    dataset_dict = json.load(file)
+
+
+# Recursively delete all entries with the specified key
+def delete_key_recursively(obj, key):
+    if isinstance(obj, dict):
+        # Remove the key if it exists in the current dictionary
+        obj.pop(key, None)
+        # Recurse into the dictionary values
+        for value in obj.values():
+            delete_key_recursively(value, key)
+    elif isinstance(obj, list):
+        # Recurse into each item in the list
+        for item in obj:
+            delete_key_recursively(item, key)
+
+
+# Call the recursive function on the loaded JSON data
+delete_key_recursively(dataset_dict, key_to_delete)
+
+# Save the modified JSON back to a file
+with open(input_path, "w") as file:
+    json.dump(dataset_dict, file, indent=4)
+
+print(f"Entries with key '{key_to_delete}' have been deleted.")

examples/registration_test.py

@@ -0,0 +1,59 @@
+import os
+import z5py
+import numpy as np
+
+from matchmaker.prealignment import prealign_sample
+from matchmaker.transform_utils import rotate_img
+from matchmaker.vis import plot_three_slices, plot_overlay
+
+
+def main():
+    output_dir = "./data/test"
+    if not os.path.exists(f"{output_dir}/plots"):
+        os.makedirs(f"{output_dir}/plots")
+
+    #######################
+    # prealign moving image
+    moving_input = "./data/platy1_muscles_stardist_moving.n5"
+    with z5py.File(moving_input, "r") as f:
+        seg_moving = f["seg"][:]
+
+    plot_three_slices(seg_moving, save_path=f"{output_dir}/plots/moving.png")
+
+    seg_moving_prealigned = prealign_sample(seg_moving, file_name="moving", save_path=output_dir)
+
+    with z5py.File(f"{output_dir}/moving_prealigned.n5", "w") as f:
+        f.create_dataset("seg", data=seg_moving_prealigned, compression="gzip")
+
+    #######################
+    # prealign fixed image
+    fixed_input = "./data/platy1_muscles_stardist_fixed.n5"
+    with z5py.File(fixed_input, "r") as f:
+        seg_fixed = f["seg"][:]
+
+    plot_three_slices(seg_fixed, save_path=f"{output_dir}/plots/fixed.png")
+
+    seg_fixed_prealigned = prealign_sample(seg_fixed, file_name="fixed", save_path=output_dir)
+
+    with z5py.File(f"{output_dir}/fixed_prealigned.n5", "w") as f:
+        f.create_dataset("seg", data=seg_fixed_prealigned, compression="gzip")
+
+    #######################
+    # test backtransform
+    T = np.loadtxt(f"{output_dir}/fixed_T_prealignment.txt")
+    seg_fixed_inv = rotate_img(
+        seg_fixed_prealigned, np.linalg.inv(T), output_shape=seg_fixed.shape
+    )
+
+    plot_three_slices(seg_fixed_inv, save_path=f"{output_dir}/plots/fixed_prealigned_inv.png")
+
+    # plot overlay of prealigned volumes
+    plot_overlay(
+        seg_fixed_prealigned,
+        seg_moving_prealigned,
+        save_path=f"{output_dir}/plots/overlay_prealignment.png",
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/reverse_deformation.py

@@ -0,0 +1,31 @@
+import numpy as np
+from matchmaker.vis import plot_three_slices, plot_overlay
+from matchmaker.transform_utils import rotate_img
+from matchmaker.n5_utils import read_volume
+
+
+def main():
+    '''
+    Reverse the deformation of a sample by applying the inverse transformation matrix.
+    '''
+    seg_moving = read_volume(
+        f="./data/platy1_muscles_stardist_moving.n5",
+        key="seg",
+    )
+
+    # compare with original image
+    seg_fixed = read_volume(
+        f="./data/platy1_muscles_stardist_fixed.n5",
+        key="seg",
+    )
+
+    T = np.loadtxt("./data/transformation_matrix.txt")
+    seg_moving_reverse = rotate_img(seg_moving, np.linalg.inv(T), output_shape=seg_fixed.shape)
+
+    plot_three_slices(seg_moving_reverse)
+    plot_three_slices(seg_fixed)
+    plot_overlay(seg_fixed, seg_moving_reverse)
+
+
+if __name__ == "__main__":
+    main()

examples/view_data.py

@@ -0,0 +1,11 @@
+from matchmaker.n5_utils import read_volume
+import napari
+
+seg_fixed = read_volume("./data/test/platy1_muscles_stardist_fixed_prealigned.n5", key="seg")
+seg_moving = read_volume("./data/test/platy1_muscles_stardist_moving_prealigned.n5", key="seg")
+# seg_moving = seg_moving[:, :, ::-1]
+
+v = napari.Viewer()
+v.add_labels(seg_fixed)
+v.add_labels(seg_moving)
+napari.run()