Register 2D and 3D (whole brain) images

brainglobe_registration/elastix/register.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import dask.array as da
 import itk
 import numpy as np
 import numpy.typing as npt
 
+from brainglobe_registration.utils.preprocess import filter_image
 from brainglobe_registration.utils.utils import (
     convert_atlas_labels,
     restore_atlas_labels,
@@ -16,7 +18,8 @@ def run_registration(
     moving_image: npt.NDArray,
     parameter_lists: List[Tuple[str, dict]],
     output_directory: Optional[Path] = None,
-) -> Tuple[npt.NDArray, itk.ParameterObject]:
+    filter_images: bool = True,
+) -> itk.ParameterObject:
     """
     Run the registration process on the given images.
 
@@ -29,18 +32,22 @@ def run_registration(
     parameter_lists : List[tuple[str, dict]]
         The list of registration parameters, one for each transform.
     output_directory : Optional[Path], optional
-        The output directory for the registration results, by default None
+        The output directory for the registration results, by default None.
+    filter_images : bool, optional
+        Whether to filter the images before registration, by default True.
 
     Returns
     -------
-    npt.NDArray
-        The result image.
     itk.ParameterObject
         The result transform parameters.
     """
+    if filter_images:
+        atlas_image = filter_image(atlas_image)
+        moving_image = filter_image(moving_image)
+
     # convert to ITK, view only
-    atlas_image = itk.GetImageViewFromArray(atlas_image).astype(itk.F)
-    moving_image = itk.GetImageViewFromArray(moving_image).astype(itk.F)
+    atlas_image = itk.GetImageViewFromArray(atlas_image)
+    moving_image = itk.GetImageViewFromArray(moving_image)
 
     # This syntax needed for 3D images
     elastix_object = itk.ElastixRegistrationMethod.New(
@@ -53,7 +60,6 @@ def run_registration(
     elastix_object.UpdateLargestPossibleRegion()
 
     # get results
-    result_image = elastix_object.GetOutput()
     result_transform_parameters = elastix_object.GetTransformParameterObject()
 
     if output_directory:
@@ -66,10 +72,7 @@ def run_registration(
             result_transform_parameters, file_names
         )
 
-    return (
-        np.asarray(result_image),
-        result_transform_parameters,
-    )
+    return result_transform_parameters
 
 
 def transform_annotation_image(
@@ -96,9 +99,15 @@ def transform_annotation_image(
     """
     adjusted_annotation_image, mapping = convert_atlas_labels(annotation_image)
 
-    annotation_image = itk.GetImageFromArray(adjusted_annotation_image).astype(
-        itk.F
-    )
+    if adjusted_annotation_image.ndim == 2:
+        adjusted_annotation_image = adjusted_annotation_image.astype(
+            np.float32
+        )
+
+    annotation_image = itk.GetImageViewFromArray(adjusted_annotation_image)
+
+    del adjusted_annotation_image
+
     temp_interp_order = transform_parameters.GetParameter(
         0, "FinalBSplineInterpolationOrder"
     )
@@ -113,8 +122,10 @@ def transform_annotation_image(
     transform_parameters.SetParameter(
         "FinalBSplineInterpolationOrder", temp_interp_order
     )
+    del annotation_image
+
     transformed_annotation_array = np.asarray(transformed_annotation).astype(
-        np.uint32
+        np.uint16
     )
 
     transformed_annotation_array = restore_atlas_labels(
@@ -143,15 +154,18 @@ def transform_image(
     npt.NDArray
         The transformed image.
     """
-    image = itk.GetImageViewFromArray(image).astype(itk.F)
+    image = itk.GetImageViewFromArray(image)
 
     transformix_object = itk.TransformixFilter.New(image)
     transformix_object.SetTransformParameterObject(transform_parameters)
     transformix_object.UpdateLargestPossibleRegion()
 
     transformed_image = transformix_object.GetOutput()
+    transformed_image = da.from_array(
+        np.asarray(transformed_image).astype(np.int16)
+    )
 
-    return np.asarray(transformed_image)
+    return transformed_image
 
 
 def calculate_deformation_field(
@@ -178,21 +192,21 @@ def calculate_deformation_field(
         The deformation field.
     """
     transformix_object = itk.TransformixFilter.New(
-        itk.GetImageViewFromArray(moving_image).astype(itk.F),
+        itk.GetImageViewFromArray(moving_image),
         transform_parameters,
     )
     transformix_object.SetComputeDeformationField(True)
 
     transformix_object.UpdateLargestPossibleRegion()
 
     # Change from ITK to numpy axes ordering
-    deformation_field = itk.GetArrayFromImage(
+    deformation_field = itk.GetArrayViewFromImage(
         transformix_object.GetOutputDeformationField()
-    )[..., ::-1]
+    )
 
     if not debug:
         # Cleanup files generated by elastix
-        (Path.cwd() / "DeformationField.tiff").unlink(missing_ok=True)
+        (Path.cwd() / "deformationField.tiff").unlink(missing_ok=True)
 
     return deformation_field
 
@@ -202,25 +216,50 @@ def invert_transformation(
     parameter_list: List[Tuple[str, dict]],
     transform_parameters: itk.ParameterObject,
     output_directory: Optional[Path] = None,
+    filter_images: bool = True,
 ) -> itk.ParameterObject:
+    """
+    Invert the transformation of the fixed image using the given transform
+    parameters.
+
+    Inverts the transformation by applying the forward transformation to the
+    fixed image and registering it to itself.
 
-    fixed_image = itk.GetImageFromArray(fixed_image).astype(itk.F)
+    Parameters
+    ----------
+    fixed_image : npt.NDArray
+        The reference image.
+    parameter_list : List[Tuple[str, dict]]
+        The list of registration parameters, one for each transform.
+    transform_parameters : itk.ParameterObject
+        The transform parameters to inverse.
+    output_directory : Optional[Path], optional
+        The output directory for the registration results, by default None.
+    filter_images : bool, optional
+        Whether to filter the images before registration, by default True.
+
+    Returns
+    -------
+    itk.ParameterObject
+        The inverse transform parameters.
+    """
+    if filter_images:
+        fixed_image = filter_image(fixed_image)
+
+    fixed_image = itk.GetImageViewFromArray(fixed_image)
 
     elastix_object = itk.ElastixRegistrationMethod.New(
         fixed_image, fixed_image
     )
 
     parameter_object_inverse = setup_parameter_object(parameter_list)
-
     elastix_object.SetInitialTransformParameterObject(transform_parameters)
-
     elastix_object.SetParameterObject(parameter_object_inverse)
 
     elastix_object.UpdateLargestPossibleRegion()
 
     num_initial_transforms = transform_parameters.GetNumberOfParameterMaps()
 
-    result_image = elastix_object.GetOutput()
     out_parameters = elastix_object.GetTransformParameterObject()
     result_transform_parameters = itk.ParameterObject.New()
 
@@ -245,10 +284,7 @@ def invert_transformation(
             result_transform_parameters, file_names
         )
 
-    return (
-        np.asarray(result_image),
-        result_transform_parameters,
-    )
+    return result_transform_parameters
 
 
 def setup_parameter_object(parameter_lists: List[tuple[str, dict]]):

brainglobe_registration/napari.yaml

@@ -5,13 +5,19 @@ contributions:
     - id: brainglobe-registration.make_registration_widget
       python_name: brainglobe_registration.registration_widget:RegistrationWidget
       title: BrainGlobe Registration
-    - id: brainglobe-registration.load_sample
-      python_name: brainglobe_registration.sample_data:load_sample_data
-      title: Sample Data
+    - id: brainglobe-registration.load_sample_2d
+      python_name: brainglobe_registration.sample_data:load_sample_data_2d
+      title: Sample Data 2D
+    - id: brainglobe-registration.load_sample_3d
+      python_name: brainglobe_registration.sample_data:load_sample_data_3d
+      title: Sample Data 3D
   sample_data:
-    - command: brainglobe-registration.load_sample
+    - command: brainglobe-registration.load_sample_2d
       display_name: 2D coronal mouse brain section
-      key: example
+      key: example_2d
+    - command: brainglobe-registration.load_sample_3d
+      display_name: 3D mouse brain
+      key: example_3d
   widgets:
     - command: brainglobe-registration.make_registration_widget
       display_name: BrainGlobe Registration

brainglobe_registration/parameters/brainglobe_registration/affine.txt

@@ -0,0 +1,56 @@
+//Affine Transformation - brainglobe-registration
+
+//ImageTypes
+(FixedInternalImagePixelType "short")
+(MovingInternalImagePixelType "short")
+
+//Components
+(Registration "MultiResolutionRegistration")
+(FixedImagePyramid "FixedRecursiveImagePyramid")
+(MovingImagePyramid "MovingRecursiveImagePyramid")
+(Interpolator "LinearInterpolator")
+(Metric "AdvancedMattesMutualInformation")
+(Optimizer "AdaptiveStochasticGradientDescent")
+(ResampleInterpolator "FinalBSplineInterpolator")
+(Resampler "DefaultResampler")
+(Transform "AffineTransform")
+
+(ErodeMask "false" )
+
+(NumberOfResolutions 3)
+
+(HowToCombineTransforms "Compose")
+(AutomaticTransformInitialization "true")
+(AutomaticScalesEstimation "true")
+
+(WriteTransformParametersEachIteration "false")
+(WriteResultImage "true")
+(ResultImageFormat "tiff")
+(CompressResultImage "false")
+(WriteResultImageAfterEachResolution "false")
+(ShowExactMetricValue "false")
+
+//Maximum number of iterations in each resolution level:
+(MaximumNumberOfIterations 500 )
+
+//Number of grey level bins in each resolution level:
+(NumberOfHistogramBins 32 )
+(FixedLimitRangeRatio 0.0)
+(MovingLimitRangeRatio 0.0)
+(FixedKernelBSplineOrder 3)
+(MovingKernelBSplineOrder 3)
+
+//Number of spatial samples used to compute the mutual information in each resolution level:
+(ImageSampler "RandomCoordinate")
+(UseRandomSampleRegion "false")
+(NumberOfSpatialSamples 3000 )
+(NewSamplesEveryIteration "true")
+(CheckNumberOfSamples "true")
+(MaximumNumberOfSamplingAttempts 10)
+
+
+//Order of B-Spline interpolation used for applying the final deformation:
+(FinalBSplineInterpolationOrder 3)
+
+//Default pixel value for pixels that come from outside the picture:
+(DefaultPixelValue 0)

brainglobe_registration/parameters/brainglobe_registration/bspline.txt

@@ -0,0 +1,59 @@
+//Bspline Transformation - brainglobe-registration
+
+//ImageTypes
+(FixedInternalImagePixelType "short")
+(MovingInternalImagePixelType "short")
+
+//Components
+(Registration "MultiResolutionRegistration")
+(FixedImagePyramid "FixedRecursiveImagePyramid")
+(MovingImagePyramid "MovingRecursiveImagePyramid")
+(Interpolator "LinearInterpolator")
+(Metric "AdvancedMattesMutualInformation")
+(Optimizer "AdaptiveStochasticGradientDescent")
+(ResampleInterpolator "FinalBSplineInterpolator")
+(Resampler "DefaultResampler")
+(Transform "RecursiveBSplineTransform")
+
+(ErodeMask "false" )
+
+(NumberOfResolutions 3)
+(FinalGridSpacingInVoxels 50.000000 50.000000 50.000000)
+
+(HowToCombineTransforms "Compose")
+
+(WriteTransformParametersEachIteration "false")
+(ResultImageFormat "tiff")
+(WriteResultImage "true")
+(CompressResultImage "false")
+(WriteResultImageAfterEachResolution "false")
+(ShowExactMetricValue "false")
+
+// Option supported in elastix 4.1:
+(UseFastAndLowMemoryVersion "true")
+
+//Maximum number of iterations in each resolution level:
+(MaximumNumberOfIterations 500 )
+
+//Number of grey level bins in each resolution level:
+(NumberOfHistogramBins 32 )
+(FixedLimitRangeRatio 0.0)
+(MovingLimitRangeRatio 0.0)
+(FixedKernelBSplineOrder 3)
+(MovingKernelBSplineOrder 3)
+
+//Number of spatial samples used to compute the mutual information in each resolution level:
+(ImageSampler "RandomCoordinate")
+(FixedImageBSplineInterpolationOrder 1 )
+(UseRandomSampleRegion "false")
+(SampleRegionSize 100.0 100.0 100.0)
+(NumberOfSpatialSamples 10000 )
+(NewSamplesEveryIteration "true")
+(CheckNumberOfSamples "true")
+(MaximumNumberOfSamplingAttempts 10)
+
+//Order of B-Spline interpolation used for applying the final deformation:
+(FinalBSplineInterpolationOrder 3)
+
+//Default pixel value for pixels that come from outside the picture:
+(DefaultPixelValue 0)