From f8fe58e8b496defc5386482794af2d803a620085 Mon Sep 17 00:00:00 2001 From: Tom Birdsong Date: Wed, 18 Oct 2023 10:18:17 -0400 Subject: [PATCH 1/2] ENH: Refactor `itk-dreg` abstract interface for plugin methods Refactors the proposed `itk-dreg` platform interface to better support extended registration methods within a `dask`-based framework: - Moves interfaces into `import`-able `itk_dreg` Python module - Restructures base types and abstract methods - Introduces helper types for working with dask blocks with metadata - Proposes updates to "register` interface for block-based registration - Proposes abstract interfaces for ITK image stream reader construction and block registration postprocessing. --- .gitignore | 2 + registration_interface.py | 46 ------ src/itk_dreg/__init__.py | 1 + src/itk_dreg/base/__init__.py | 0 src/itk_dreg/base/image_block_interface.py | 151 +++++++++++++++++++ src/itk_dreg/base/itk_typing.py | 28 ++++ src/itk_dreg/base/registration_interface.py | 158 ++++++++++++++++++++ 7 files changed, 340 insertions(+), 46 deletions(-) create mode 100644 .gitignore delete mode 100644 registration_interface.py create mode 100644 src/itk_dreg/__init__.py create mode 100644 src/itk_dreg/base/__init__.py create mode 100644 src/itk_dreg/base/image_block_interface.py create mode 100644 src/itk_dreg/base/itk_typing.py create mode 100644 src/itk_dreg/base/registration_interface.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..372c13e --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ + diff --git a/registration_interface.py b/registration_interface.py deleted file mode 100644 index 861c047..0000000 --- a/registration_interface.py +++ /dev/null @@ -1,46 +0,0 @@ -import itk -from abc import ABC, abstractmethod - -# The function signature each registration approach should provide - -FloatImage2DType = itk.Image[itk.F, 2] -FloatImage3DType = itk.Image[itk.F, 3] -FloatImageType = Union[Image2DType, Image3DType] - -class RegistrationMethod(ABC): - @abstractmethod - def __call__( - self, - fixed_image: FloatImageType, - moving_image: FloatImageType, - initial_transform: itk.Transform[itk.D, 3], - **kwargs - ) -> itk.Transform[itk.D, 3]: - pass - - -""" -my_moving_image = ... -my_fixed_image = ... - -my_initial_transform = ... - -my_transform = registration_method(my_fixed_image, my_moving_image, my_initial_transform) - -# registration method returns an update to the initial transform - -final_transform = itk.CompositeTransform() -final_transform.append_transform(my_initial_transform) -final_transform.append_transform(my_transform) - -interpolator = itk.LinearInterpolateImageFunction.New(my_moving_image) - -my_warped_image = itk.resample_image_filter( - my_moving_image, - transform=final_transform, - interpolator=interpolator, - use_reference_image=True, - reference_image=my_fixed_image -) - -""" diff --git a/src/itk_dreg/__init__.py b/src/itk_dreg/__init__.py new file mode 100644 index 0000000..f102a9c --- /dev/null +++ b/src/itk_dreg/__init__.py @@ -0,0 +1 @@ +__version__ = "0.0.1" diff --git a/src/itk_dreg/base/__init__.py b/src/itk_dreg/base/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/itk_dreg/base/image_block_interface.py b/src/itk_dreg/base/image_block_interface.py new file mode 100644 index 0000000..3427fd7 --- /dev/null +++ b/src/itk_dreg/base/image_block_interface.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 + +from dataclasses import dataclass +from enum import IntEnum +from typing import Optional, List + +import dask.array +import numpy.typing as npt + +from .itk_typing import ImageType, TransformType + +""" +Define common data structures for managing block regions and registration output. +""" + + +class BlockRegStatus(IntEnum): + """ + Status codes indicating the registration outcome from a single block pair. + + TODO: To be extended with more granular error codes for `itk-dreg` infrastructure. + """ + + SUCCESS = 0 + """Registration yielded at least a forward transform result.""" + + FAILURE = 1 + """Registration encountered an unspecified error.""" + + +@dataclass +class BlockInfo: + """ + Header information describing the position of a lazy dask subvolume (block) + in voxel space with respect to a parent volume. + + Accessors are in NumPy order: (K,J,I) where K is slowest, I is fastest + """ + + chunk_index: List[int] + """ + The chunk position in the greater image volume in terms of chunk counts. + For instance, if traversing 2x2x2 volume along the fastest axis first: + - the 0th chunk would have chunk index (0,0,0), + - the 1st chunk would have chunk index (0,0,1), + - the 7th chunk would have chunk index (1,0,0) + """ + + array_slice: List[slice] + """ + The chunk position in the greater image volume in terms of voxel access. + For instance, if a 100x100x100 volume is evenly subdivided into 10x10x10 chunks, + the first chunk would slice along [(0,10,1),(0,10,1),(0,10,1)]. + """ + + +@dataclass +class LocatedBlock: + """ + Combined header and data access to get a lazy dask volume with respect + to a parent volume in voxel space. + + Accessors are in NumPy order: (K,J,I) where K is slowest, I is fastest + """ + + loc: BlockInfo + """ + The location of the block relative to the parent image voxel array. + """ + + arr: dask.array.core.Array + """ + The dask volume for lazy voxel access. + """ + + +@dataclass +class BlockPairRegistrationResult: + """Encapsulate result of fixed-to-moving registration over one block pair.""" + + transform: Optional[TransformType] + """ + The forward transform registration result, if any. + The forward transform maps from moving to fixed space. + """ + + transform_domain: Optional[ImageType] + """ + Oriented representation of the domain over which the forward transform is valid. + `transform_domain` has no voxel data and serves as a metadata representation of an + oriented bounding box in physical space. + `transform_domain` must be available if and only if `transform` is available. + """ + + inv_transform: Optional[TransformType] + """ + The inverse transform registration result, if any. + The inverse transform maps from fixed to moving space. + If `inv_transform` is available then `transform` must also be available. + """ + + inv_transform_domain: Optional[ImageType] + """ + Oriented representation of the domain over which the inverse transform is valid. + `inv_transform_domain` has no voxel data and serves as a metadata representation of an + oriented bounding box in physical space. + `inv_transform_domain` must be available if and only if `inv_transform` is available. + """ + + status: BlockRegStatus + """Status code indicating registration success or failure.""" + + +@dataclass +class RegistrationTransformResult: + """ + Encapsulate result of fixed-to-moving registration over all block pairs. + """ + + transform: TransformType + """ + The forward transform resulting from block postprocessing. + The forward transform maps from moving to fixed space. + """ + + inv_transform: Optional[TransformType] + """ + The inverse transform registration result from block postprocessing, if any. + The inverse transform maps from fixed to moving space. + If `inv_transform` is available then `transform` must also be available. + """ + + +@dataclass +class RegistrationResult: + """ + Encapsulate result of fixed-to-moving registration over all block pairs. + """ + + transforms: RegistrationTransformResult + """The forward and inverse transforms resulting from registration.""" + + status: npt.ArrayLike + """ + `status` is an ND array where each element reflects the status code output + for block pair registration over the corresponding input moving chunk. + + `status` has the same shape as the moving input array of chunks. + That is, if the moving input array is subdivided into 2 chunks x 3 chunks x 4 chunks, + `status` will be an array of voxels with shape [2,3,4]. + """ diff --git a/src/itk_dreg/base/itk_typing.py b/src/itk_dreg/base/itk_typing.py new file mode 100644 index 0000000..952fb60 --- /dev/null +++ b/src/itk_dreg/base/itk_typing.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 + +from typing import Union + +import itk + +""" +Define common "union" type hints for floating-point registration +in two or three dimensional space. +""" + +ImagePixelType = itk.F +FloatImage2DType = itk.Image[ImagePixelType, 2] +FloatImage3DType = itk.Image[ImagePixelType, 3] +ImageType = Union[FloatImage2DType, FloatImage3DType] + +ImageRegion2DType = itk.ImageRegion[2] +ImageRegion3DType = itk.ImageRegion[3] +ImageRegionType = Union[ImageRegion2DType, ImageRegion3DType] + +FloatImage2DReaderType = itk.ImageFileReader[FloatImage2DType] +FloatImage3DReaderType = itk.ImageFileReader[FloatImage3DType] +ImageReaderType = Union[FloatImage2DReaderType, FloatImage3DReaderType] + +TransformScalarType = itk.D +TransformType = Union[ + itk.Transform[TransformScalarType, 2], itk.Transform[TransformScalarType, 3] +] diff --git a/src/itk_dreg/base/registration_interface.py b/src/itk_dreg/base/registration_interface.py new file mode 100644 index 0000000..ffc9e6a --- /dev/null +++ b/src/itk_dreg/base/registration_interface.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 + +import itk +from abc import ABC, abstractmethod +from typing import Optional, Iterable, Tuple + +from .image_block_interface import ( + BlockPairRegistrationResult, + RegistrationTransformResult, +) +from .itk_typing import ImageType, ImageReaderType, ImageRegionType, TransformType + +""" +Defines extensible components to extend with concrete implementations. +""" + + +class ConstructReaderMethod(ABC): + """ + A method that generates a new `itk.ImageFileReader` for image registration. + + ITK provides the `itk.ImageFileReader` mechanism to retrieve all or part of + a spatial image from a provided local or remote image source. `itk-dreg` + registration infrastructure attempts to stream image subregions into memory + at runtime in order to perform block-based pairwise registration without + ever loading an entire image into memory at once. + + Extend this class to customize the image reading step for cases such as + to attach domain-specific metadata or to convert from a reference type + not supported by ITK by default. + + The resulting `itk.ImageFileReader` object MUST be initialized with metadata + to represent the extent of the underlying data in voxel and physical space. + + It is strongly recommended that the resulting `itk.ImageFileReader` is NOT + initialized with underlying voxel data. Voxel regions should be lazily + initialized by `itk-dreg` registration infrastructure to match block + requested regions. + + .. code-block:: python + + ReaderSource = ConstructReaderMethodSubclass(...) + image_source = ReaderSource() + image = image_source.UpdateLargestPossibleRegion() + """ + + @abstractmethod + def __call__(self, **kwargs) -> ImageReaderType: + pass + + +class BlockPairRegistrationMethod(ABC): + """ + A method that registers two spatially located image blocks together. + + `fixed_subimage` and `moving_image` inputs are `itk.Image` representations + of block subregions within greater fixed and moving inputs. + + Extend this class to implement custom registration method that plugs in + to `itk-dreg` registration infrastructure. + """ + + @abstractmethod + def __call__( + self, + fixed_subimage: ImageType, + moving_subimage: ImageType, + initial_transform: TransformType, + **kwargs + ) -> BlockPairRegistrationResult: + """ + Run image-to-image pairwise registration. + + :param fixed_subimage: The reference fixed subimage. + `fixed_subimage.RequestedRegion` reflects the requested subregion corresponding + to the scheduled dask array chunk. + The initial `fixed_subimage.BufferedRegion` includes the requested region + and possibly an extra padding factor introduced before fetching fixed + image voxel data. + :param moving_subimage: The moving subimage to be registered onto fixed image space. + `moving_subimage.RequestedRegion` reflects the requested subregion + corresponding to the approximate physical bounds of `fixed_subimage.RequestedRegion` + after initialization with `initial_transform`. + The initial `moving_subimage.BufferedRegion` includes the requested region + and possibly an extra padding factor introduced before fetching fixed + image voxel data. + :param initial_transform: The forward transform representing an initial alignment + mapping from fixed to moving image space. + """ + pass + + +class ReduceResultsMethod(ABC): + """ + A method that reduces a sparse collection of pairwise block registration results + to yield a generalized fixed-to-moving transform. + + Extend this class to implement a custom method mapping block results to a general transform. + + Possible implementations could include methods for finding global consensus among results, + stitching methods to yield a piecewise transform, or patchwise methods to normalize among + bounded transform domains. + """ + + @abstractmethod + def __call__( + self, + block_results: Iterable[Tuple[ImageType, BlockPairRegistrationResult]], + fixed_reader_ctor: ConstructReaderMethod, + initial_transform: itk.Transform, + **kwargs + ) -> RegistrationTransformResult: + """ + :param block_results: An iterable collection of subimages in fixed space + and the corresponding registration result for the given subimage. + Subimages are not buffered and represent the subdomains within the + original fixed image space prior to initial transform application. + :param fixed_reader_ctor: Method to create an image reader to stream + part or all of the fixed image. + :param initial_transform: Initial forward transform used in registration. + The forward transform maps from the fixed to moving image. + """ + pass + + +""" +my_fixed_image = ... +my_moving_image = ... + +my_initial_transform = ... + +# registration method returns an update to the initial transform + +my_transform = itk.register_dreg( + construct_fixed_image_method=my_construct_streaming_reader_method, + construct_moving_image_method=my_construct_streaming_reader_method, + initial_transform=my_initial_transform, + registration_method=my_block_pair_registration_method_subclass, + reduce_method=my_reduce_registration_method_subclass +) + +final_transform = itk.CompositeTransform() +final_transform.append_transform(my_initial_transform) +final_transform.append_transform(my_transform) + +# we can use the result transform to resample the moving image to fixed image space + +interpolator = itk.LinearInterpolateImageFunction.New(my_moving_image) + +my_warped_image = itk.resample_image_filter( + my_moving_image, + transform=final_transform, + interpolator=interpolator, + use_reference_image=True, + reference_image=my_fixed_image +) + +""" From 523293f0c51c50b6d0130efcb154e541dcaa732f Mon Sep 17 00:00:00 2001 From: Tom Birdsong Date: Wed, 18 Oct 2023 10:19:35 -0400 Subject: [PATCH 2/2] ENH: Add pytest to validate imports Add simple `pytest` infrastructure to validate that `itk_dreg` module scripts load without error. --- test/test_import.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 test/test_import.py diff --git a/test/test_import.py b/test/test_import.py new file mode 100644 index 0000000..c3d4f64 --- /dev/null +++ b/test/test_import.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 + +# Purpose: Simple pytest to validate that `itk-dreg` modules can be loaded. + +import sys + +sys.path.append("src") + +import itk + +itk.auto_progress(2) + + +def test_loadmodule(): + import itk_dreg.base.image_block_interface + import itk_dreg.base.itk_typing + import itk_dreg.base.registration_interface