adding affine registration

code/aind_ccf_reg/register.py

@@ -236,7 +236,12 @@ def _qc_reg(
         if moved_path:
             self._plot_write_antsimg(ants_moved, moved_path)
 
-    def register_to_template(self, ants_fixed, ants_moving):
+    def register_to_template(
+        self,
+        ants_fixed,
+        ants_moving,
+        moving_mask=None,
+    ):
         """
         Run SyN regsitration to align brain image to SPIM template
 
@@ -246,6 +251,8 @@ def register_to_template(self, ants_fixed, ants_moving):
             fixed image
         ants_moving: ANTsImage
             moving image
+        moving_mask: ANTsImage
+            Moving mask
 
         Returns
         -----------
@@ -264,12 +271,22 @@ def register_to_template(self, ants_fixed, ants_moving):
         registration_params = {
             "fixed": ants_fixed,
             "moving": ants_moving,
+            "moving_mask": moving_mask,
             "type_of_transform": "Rigid",
             "outprefix": f"{self.args['results_folder']}/ls_to_template_rigid_",
             "mask_all_stages": True,
             "grad_step": 0.25,
-            "reg_iterations": (60, 40, 20, 0),
+            "reg_iterations": [
+                0,
+                0,
+                0,
+                0,
+            ],  # Explicitly avoiding deformable reg
+            "reg_iterations": [60, 30, 15, 5],
             "aff_metric": "mattes",
+            "verbose": True,  # Setting to true for future debugging
+            "flow_sigma": 3.0,
+            "total_sigma": 0.0,
         }
 
         logger.info(
@@ -292,14 +309,55 @@ def register_to_template(self, ants_fixed, ants_moving):
             figpath_name=reg_task,
         )
 
+        # ----------------------------------#
+        # Affine registration
+        # ----------------------------------#
+        logger.info(f"Start computing affine registration ....")
+
+        start_time = datetime.now()
+        affine_registration_params = {
+            "fixed": ants_fixed,
+            "moving": ants_moving,  # The moving image after rigid registration
+            "moving_mask": moving_mask,
+            "initial_transform": [
+                f"{self.args['results_folder']}/ls_to_template_rigid_0GenericAffine.mat"
+            ],
+            "outprefix": f"{self.args['results_folder']}/ls_to_template_affine_",
+            "type_of_transform": "Affine",
+            "reg_iterations": [0, 0, 0, 0],
+            "aff_iterations": [60, 30, 15, 5],
+            "aff_metric": "mattes",  # CC metric?
+            "verbose": True,
+            "mask_all_stages": True,
+        }
+        logger.info(
+            f"Computing affine registration with parameters: {affine_registration_params}"
+        )
+        affine_reg = ants.registration(**affine_registration_params)
+
+        end_time = datetime.now()
+        logger.info(
+            f"Affine registration Complete, execution time: {end_time - start_time} s -- image {affine_reg}"
+        )
+        ants_moved = affine_reg["warpedmovout"]
+
+        reg_task = "reg_affine"
+        self._qc_reg(
+            ants_moving,
+            ants_fixed,
+            ants_moved,
+            moved_path=self.args["ants_params"].get("affine_path"),
+            figpath_name=reg_task,
+        )
+
         # ----------------------------------#
         # SyN registration
         # ----------------------------------#
 
         logger.info(f"Start registering to template ....")
 
         if self.args["reference_res"] == 25:
-            reg_iterations = [200, 20, 0]
+            reg_iterations = [200, 100, 25, 3]
         elif self.args["reference_res"] == 10:
             reg_iterations = [400, 200, 40, 0]
         else:
@@ -311,12 +369,16 @@ def register_to_template(self, ants_fixed, ants_moving):
         registration_params = {
             "fixed": ants_fixed,
             "moving": ants_moving,
-            # "initial_transform": [f"{self.args['reg_folder']}/ls_to_template_rigid_0GenericAffine.mat"],
-            "initial_transform": rigid_reg["fwdtransforms"][0],
+            "moving_mask": moving_mask,
+            "reg_iterations": reg_iterations,
+            "initial_transform": [
+                f"{self.args['results_folder']}/ls_to_template_affine_0GenericAffine.mat"
+            ],
+            "mask_all_stages": True,
             "syn_metric": "CC",
             "syn_sampling": 2,
-            "reg_iterations": reg_iterations,
             "outprefix": f"{self.args['results_folder']}/ls_to_template_SyN_",
+            "verbose": True,
         }
 
         logger.info(

code/main.py

@@ -16,9 +16,7 @@ def main() -> None:
     """
     data_folder = os.path.abspath("../data")
     processing_manifest_path = f"{data_folder}/processing_manifest.json"
-    acquisition_path = (
-        f"{data_folder}/acquisition.json"
-    )
+    acquisition_path = f"{data_folder}/acquisition.json"
 
     if not os.path.exists(processing_manifest_path):
         raise ValueError("Processing manifest path does not exist!")
@@ -171,6 +169,7 @@ def main() -> None:
                 "right_to_left": 0,
             },
             "rigid_path": f"{reg_folder}/moved_rigid.nii.gz",
+            "affine_path": f"{reg_folder}/moved_affine.nii.gz",
             "moved_to_template_path": f"{reg_folder}/moved_ls_to_template.nii.gz",
             "moved_to_ccf_path": f"{results_folder}/moved_ls_to_ccf.nii.gz",
             "ccf_anno_to_brain_path": f"{reg_folder}/moved_ccf_anno_to_ls.nii.gz",

tests/test_ccf_registration.py

@@ -2,11 +2,12 @@
 
 import os
 import unittest
+from pathlib import Path
 
 import ants
 import numpy as np
 from aind_ccf_reg.preprocess import Masking
-from pathlib import Path
+
 
 class CCFRegistrationTest(unittest.TestCase):
     """Tests CCF registration."""