Merge pull request #23 from AllenNeuralDynamics/fix-add-affine

adding affine registration

code/aind_ccf_reg/register.py

@@ -11,17 +11,14 @@
 import logging
 import multiprocessing
 import os
-import shutil
 from datetime import datetime
-from glob import glob
 from pathlib import Path
 from typing import Dict, Hashable, List, Sequence, Tuple, Union
 
 import ants
 import dask
 import dask.array as da
 import numpy as np
-import tifffile
 import xarray_multiscale
 import zarr
 from aicsimageio.types import PhysicalPixelSizes
@@ -31,7 +28,6 @@
 from dask.distributed import Client, LocalCluster, performance_report
 from distributed import wait
 from numcodecs import blosc
-from skimage import io
 
 from .__init__ import __version__
 
@@ -236,7 +232,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 +247,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 +267,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 +305,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",
+            "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 +365,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(
@@ -480,7 +538,7 @@ def atlas_alignment(
 
         # ants_img = ants.image_read(self.args["prep_params"].get("percNorm_path")) #
 
-        # register to SPIM template: rigid + SyN
+        # register to SPIM template: rigid + affine + SyN
         aligned_image = self.register_to_template(ants_template, ants_img)
 
         # ----------------------------------#

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",

environment/Dockerfile

@@ -25,6 +25,5 @@ RUN pip install -U --no-cache-dir \
     matplotlib==3.7.3 \
     ome-zarr==0.8.2 \
     natsort==8.4.0 \
-    aicsimageio@git+https://github.com/camilolaiton/aicsimageio.git@feature/zarrwriter-multiscales-daskjobs
-
-RUN conda install -c conda-forge awscli
+    aicsimageio@git+https://github.com/camilolaiton/aicsimageio.git@feature/zarrwriter-multiscales-daskjobs \
+    awscli

tests/test_ccf_registration.py

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