diff --git a/brainles_preprocessing/modality.py b/brainles_preprocessing/modality.py index 0f67cde..9c5e5e3 100644 --- a/brainles_preprocessing/modality.py +++ b/brainles_preprocessing/modality.py @@ -239,6 +239,7 @@ def register( registration_dir: Union[str, Path], moving_image_name: str, step: PreprocessorSteps, + interpolator: Optional[str] = None, ) -> Path: """ Register the current modality to a fixed image using the specified registrator. @@ -249,9 +250,13 @@ def register( registration_dir (str or Path): Directory to store registration results. moving_image_name (str): Name of the moving image. step (PreprocessorSteps): The current preprocessing step. + interpolator (Optional[str]): Interpolation method used during registration. Only supported by ANTsRegistrator; omit it for other registrators to use the backend default. Returns: Path: Path to the registration matrix. + + Raises: + ValueError: If an interpolator is passed for a registrator other than ANTsRegistrator. """ fixed_image_path = Path(fixed_image_path) registration_dir = Path(registration_dir) @@ -262,12 +267,21 @@ def register( # Note, add file ending depending on registration backend! registered_matrix = registration_dir / f"M_{moving_image_name}" + if interpolator is not None and not isinstance(registrator, ANTsRegistrator): + raise ValueError( + f"'interpolator' is only supported by ANTsRegistrator for registration, " + f"got {type(registrator).__name__}. Omit it to use the backend default." + ) + + interpolator_kwargs = {"interpolator": interpolator} if interpolator else {} + registrator.register( fixed_image_path=fixed_image_path, moving_image_path=self.current, transformed_image_path=registered, matrix_path=registered_matrix, log_file_path=str(registered_log), + **interpolator_kwargs, ) self.current = registered self.steps[step] = registered @@ -364,6 +378,7 @@ def transform( moving_image_name: str, transformation_matrix_path: Union[str, Path], step: PreprocessorSteps, + interpolator: Optional[str] = None, ) -> None: """ Transform the current modality using the specified registrator and transformation matrix. @@ -375,8 +390,12 @@ def transform( moving_image_name (str): Name of the moving image. transformation_matrix_path (str or Path): Path to the transformation matrix. step (PreprocessorSteps): The current preprocessing step. + interpolator (Optional[str]): Interpolation method used during transformation. "linear" and "nearestNeighbor" are supported by ANTsRegistrator, NiftyRegRegistrator and GreedyRegistrator; further options depend on the chosen registrator. ElastixRegistrator does not support choosing an interpolator. Returns: None + + Raises: + ValueError: If an interpolator is passed for a registrator that does not support it. """ fixed_image_path = Path(fixed_image_path) registration_dir_path = Path(registration_dir_path) @@ -385,6 +404,12 @@ def transform( transformed = registration_dir_path / f"{moving_image_name}.nii.gz" transformed_log = registration_dir_path / f"{moving_image_name}.log" + interpolator_kwargs = ( + {"interpolator": registrator.normalize_interpolator(interpolator)} + if interpolator + else {} + ) + if ( isinstance(registrator, (ANTsRegistrator, NiftyRegRegistrator)) and step == PreprocessorSteps.ATLAS_REGISTERED @@ -406,6 +431,7 @@ def transform( transformation_matrix_path, # atlas registration matrix ], log_file_path=str(transformed_log), + **interpolator_kwargs, ) else: registrator.transform( @@ -414,6 +440,7 @@ def transform( transformed_image_path=str(transformed), matrix_path=str(transformation_matrix_path), log_file_path=str(transformed_log), + **interpolator_kwargs, ) self.current = transformed @@ -707,6 +734,7 @@ def deface( transformed_image_path=mask_path, matrix_path=atlas_bet_M, log_file_path=defaced_dir_path / "inverse_transform.log", + interpolator="nearestNeighbor", ) else: defacer.deface( diff --git a/brainles_preprocessing/registration/ANTs/ANTs.py b/brainles_preprocessing/registration/ANTs/ANTs.py index fa0808e..142a54c 100644 --- a/brainles_preprocessing/registration/ANTs/ANTs.py +++ b/brainles_preprocessing/registration/ANTs/ANTs.py @@ -25,6 +25,8 @@ class ANTsRegistrator(Registrator): + INTERPOLATOR_MAP = {"linear": "linear", "nearestNeighbor": "nearestNeighbor"} + def __init__( self, registration_params: Optional[Dict[str, Any]] = None, @@ -61,6 +63,7 @@ def register( transformed_image_path: Union[str, Path], matrix_path: Union[str, Path], log_file_path: Union[str, Path], + interpolator: Optional[str] = None, **kwargs, ) -> None: """ @@ -72,6 +75,7 @@ def register( transformed_image_path (str or Path): Path to the transformed image (output). matrix_path (str or Path): Path to the transformation matrix (output). log_file_path (str or Path): Path to the log file. + interpolator (Optional[str]): Interpolator to use for the transformation. Defaults to the instantiated transformation parameters, or 'linear' if not set. **kwargs: Additional registration parameters to update the instantiated defaults. Raises: @@ -83,6 +87,10 @@ def register( # We update the registration parameters with the provided kwargs registration_kwargs = {**self.registration_params, **kwargs} + interpolator = self._resolve_interpolator( + interpolator, dict(self.transformation_params) + ) + # Convert all paths to Path objects fixed_image_path = Path(fixed_image_path) moving_image_path = Path(moving_image_path) @@ -123,6 +131,7 @@ def register( transformed_image_path=transformed_image_path, matrix_path=matrix_path, log_file_path=log_file_path, + interpolator=interpolator, ) end_time = datetime.datetime.now() @@ -136,6 +145,7 @@ def register( moving_image_path=moving_image_path, transformed_image_path=transformed_image_path, matrix_path=matrix_path, + interpolator=interpolator, operation_name="registration", start_time=start_time, end_time=end_time, @@ -148,7 +158,7 @@ def transform( transformed_image_path: Union[str, Path], matrix_path: str | Path | List[str | Path], log_file_path: Union[str, Path], - interpolator: str = "nearestNeighbor", + interpolator: Optional[str] = None, **kwargs, ) -> None: """ @@ -161,7 +171,7 @@ def transform( transformed_image_path (str or Path): Path to the transformed image (output). matrix_path (str or Path or List[str | Path]): Path to the transformation matrix or a list of matrices. log_file_path (str or Path): Path to the log file. - interpolator (str): Interpolator to use for the transformation. Default is 'nearestNeighbor'. + interpolator (Optional[str]): Interpolator to use for the transformation. Defaults to the instantiated transformation parameters, or 'linear' if not set. **kwargs: Additional transformation parameters to update the instantiated defaults. Raises: AssertionError: If the interpolator is not valid. @@ -169,15 +179,12 @@ def transform( """ start_time = datetime.datetime.now() - assert interpolator in VALID_INTERPOLATORS, ( - f"Invalid interpolator: {interpolator}. " - f"Valid options are: {', '.join(VALID_INTERPOLATORS)}." - ) - # TODO - self.transformation_params # we update the transformation parameters with the provided kwargs transform_kwargs = {**self.transformation_params, **kwargs} + interpolator = self._resolve_interpolator(interpolator, transform_kwargs) + # Convert all paths to Path objects fixed_image_path = Path(fixed_image_path) moving_image_path = Path(moving_image_path) @@ -243,7 +250,7 @@ def inverse_transform( transformed_image_path: Union[str, Path], matrix_path: str | Path | List[str | Path], log_file_path: Union[str, Path], - interpolator: str = "nearestNeighbor", + interpolator: Optional[str] = None, **kwargs, ) -> None: """ @@ -255,7 +262,7 @@ def inverse_transform( transformed_image_path (str or Path): Path to the transformed image (output). matrix_path (str or Path or List[str | Path]): Path to the transformation matrix or a list of matrices. log_file_path (str or Path): Path to the log file. - interpolator (str): Interpolator to use for the transformation. Default is 'nearestNeighbor'. + interpolator (Optional[str]): Interpolator to use for the transformation. Defaults to the instantiated transformation parameters, or 'linear' if not set. **kwargs: Additional transformation parameters to update the instantiated defaults. """ if not isinstance(matrix_path, list): @@ -271,6 +278,31 @@ def inverse_transform( **kwargs, ) + def _resolve_interpolator( + self, interpolator: Optional[str], transform_kwargs: Dict[str, Any] + ) -> str: + """ + Resolve the interpolator to use. The argument takes precedence over the instantiated transformation parameters, which default to 'linear'. + + Args: + interpolator (Optional[str]): Interpolator passed to the calling method. + transform_kwargs (Dict[str, Any]): Transformation parameters; the 'interpolator' key is removed. + + Returns: + str: The validated interpolator. + + Raises: + AssertionError: If the interpolator is not valid. + """ + default_interpolator = transform_kwargs.pop("interpolator", "linear") + interpolator = self.normalize_interpolator(interpolator or default_interpolator) + + assert interpolator in VALID_INTERPOLATORS, ( + f"Invalid interpolator: {interpolator}. " + f"Valid options are: {', '.join(VALID_INTERPOLATORS)}." + ) + return interpolator + @staticmethod def _log_to_file( log_file_path: Union[str, Path], diff --git a/brainles_preprocessing/registration/elastix/elastix.py b/brainles_preprocessing/registration/elastix/elastix.py index b5780df..7dd8367 100644 --- a/brainles_preprocessing/registration/elastix/elastix.py +++ b/brainles_preprocessing/registration/elastix/elastix.py @@ -80,7 +80,7 @@ def transform( transformed_image_path: str, matrix_path: str, log_file_path: Optional[str] = None, - interpolator: str = "linear", + interpolator: Optional[str] = None, **kwargs ) -> None: """ @@ -92,7 +92,14 @@ def transform( transformed_image_path (str): Path to the transformed image (output). matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists. log_file_path (Optional[str]): Path to the log file. + interpolator (Optional[str]): Not supported by elastix. Must be None. + + Raises: + ValueError: If an interpolator is provided. """ + if interpolator is not None: + self.normalize_interpolator(interpolator) + parameter_object = self.__initialize_parameter_object() # check if the matrix file exists diff --git a/brainles_preprocessing/registration/greedy/greedy.py b/brainles_preprocessing/registration/greedy/greedy.py index 9a63b60..f8ec70c 100644 --- a/brainles_preprocessing/registration/greedy/greedy.py +++ b/brainles_preprocessing/registration/greedy/greedy.py @@ -8,8 +8,12 @@ from brainles_preprocessing.registration.registrator import Registrator from brainles_preprocessing.utils import check_and_add_suffix +VALID_INTERPOLATORS = ["NN", "LINEAR", "LABEL"] + class GreedyRegistrator(Registrator): + INTERPOLATOR_MAP = {"linear": "LINEAR", "nearestNeighbor": "NN"} + def __init__( self, ): @@ -72,10 +76,15 @@ def transform( transformed_image_path (str): Path to the transformed image (output). matrix_path (str): Path to the transformation matrix (output). This gets overwritten if it already exists. log_file_path (Optional[str]): Path to the log file, which is not used. - interpolator (Optional[str]): The interpolator to use; one of NN, LINEAR or LABEL. Defaults to LINEAR. + interpolator (Optional[str]): The interpolator to use; one of NN, LINEAR or LABEL, or "linear"/"nearestNeighbor". Defaults to LINEAR. + Raises: + AssertionError: If the interpolator is not valid. """ registor = Greedy3D() - interpolator_upper = interpolator.upper() + interpolator_upper = self.normalize_interpolator(interpolator).upper() + assert ( + interpolator_upper in VALID_INTERPOLATORS + ), f"Invalid interpolator: {interpolator}. Valid options are: {', '.join(VALID_INTERPOLATORS)}" if "LABEL" in interpolator_upper: interpolator_upper += " 0.3vox" diff --git a/brainles_preprocessing/registration/niftyreg/niftyreg.py b/brainles_preprocessing/registration/niftyreg/niftyreg.py index 3be4aa5..2748cc3 100644 --- a/brainles_preprocessing/registration/niftyreg/niftyreg.py +++ b/brainles_preprocessing/registration/niftyreg/niftyreg.py @@ -21,6 +21,8 @@ class NiftyRegRegistrator(Registrator): + INTERPOLATOR_MAP = {"linear": "1", "nearestNeighbor": "0"} + def __init__( self, registration_abspath: str = os.path.dirname(os.path.abspath(__file__)), @@ -155,7 +157,7 @@ def transform( transformed_image_path: str, matrix_path: str | Path | List[str | Path], log_file_path: str, - interpolator: str = "0", + interpolator: str = "1", **kwargs: dict, ) -> None: """ @@ -168,10 +170,11 @@ def transform( transformed_image_path (str): Path to the transformed image (output). matrix_path (str | Path | List[str | Path]): Path(s) to the transformation matrix(es). log_file_path (str): Path to the log file. - interpolator (str): Interpolation order (0, 1, 3, 4) (0=NN, 1=LIN; 3=CUB, 4=SINC). Default is '1' (linear). + interpolator (str): Interpolation order (0, 1, 3, 4) (0=NN, 1=LIN; 3=CUB, 4=SINC) or "linear"/"nearestNeighbor". Default is '1' (linear). Raises: AssertionError: If the interpolator is not valid. """ + interpolator = self.normalize_interpolator(interpolator) assert ( interpolator in VALID_INTERPOLATORS ), f"Invalid interpolator: {interpolator}. Valid options are: {', '.join([f'{k} ({v})' for k, v in VALID_INTERPOLATORS.items()])}" @@ -236,7 +239,7 @@ def inverse_transform( transformed_image_path: str, matrix_path: str | Path | List[str | Path], log_file_path: str, - interpolator: str = "0", + interpolator: str = "1", ) -> None: """ Apply inverse transformation using NiftyReg. diff --git a/brainles_preprocessing/registration/registrator.py b/brainles_preprocessing/registration/registrator.py index 17fc486..1653652 100644 --- a/brainles_preprocessing/registration/registrator.py +++ b/brainles_preprocessing/registration/registrator.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Dict, Optional class Registrator(ABC): @@ -7,6 +7,30 @@ class Registrator(ABC): # def __init__(self, backend): # self.backend = backend + # Maps the shared interpolator names ("linear", "nearestNeighbor") to the backend-specific value. + # None means the backend does not support choosing an interpolator. + INTERPOLATOR_MAP: Optional[Dict[str, str]] = None + + @classmethod + def normalize_interpolator(cls, interpolator: str) -> str: + """ + Map a shared interpolator name to the backend-specific value. Backend-specific values are passed through unchanged. + + Args: + interpolator (str): Shared name ("linear", "nearestNeighbor") or backend-specific value. + + Returns: + str: The backend-specific interpolator value. + + Raises: + ValueError: If the backend does not support choosing an interpolator. + """ + if cls.INTERPOLATOR_MAP is None: + raise ValueError( + f"{cls.__name__} does not support choosing an interpolator." + ) + return cls.INTERPOLATOR_MAP.get(interpolator, interpolator) + @abstractmethod def register( self, @@ -37,7 +61,7 @@ def transform( matrix_path: Any, log_file_path: str, interpolator: str = "default", - **kwargs + **kwargs, ): """ Abstract method for transforming images. diff --git a/brainles_preprocessing/transform.py b/brainles_preprocessing/transform.py index 4447c90..56d24fb 100644 --- a/brainles_preprocessing/transform.py +++ b/brainles_preprocessing/transform.py @@ -53,9 +53,9 @@ def apply( moving_image (Union[str, Path]): Path to the moving image. E.g., this could be a segmentation in atlas space. output_img_path (Union[str, Path]): Path where the transformed image will be saved. log_file_path (Union[str, Path]): Path to the log file where transformation details will be written. - interpolator (Optional[str]): Interpolation method used during transformation. + interpolator (Optional[str]): Interpolation method used during transformation. Defaults to "nearestNeighbor", which is suited for labels/segmentations; use "linear" for intensity images. inverse (bool): If True, applies the inverse transformation. Default is False. - Available options depend on the chosen registrator: + "linear" and "nearestNeighbor" are supported by ANTsRegistrator, NiftyRegRegistrator and GreedyRegistrator. ElastixRegistrator does not support choosing an interpolator. Further options depend on the chosen registrator: - **ANTsRegistrator**: - "linear" @@ -104,6 +104,8 @@ def apply( "matrix_path": transforms, "log_file_path": str(log_file_path), } + if interpolator is None and self.registrator.INTERPOLATOR_MAP is not None: + interpolator = "nearestNeighbor" if interpolator is not None: kwargs["interpolator"] = interpolator diff --git a/tests/test_modality.py b/tests/test_modality.py new file mode 100644 index 0000000..932f708 --- /dev/null +++ b/tests/test_modality.py @@ -0,0 +1,48 @@ +import unittest +from pathlib import Path + +from brainles_preprocessing.constants import PreprocessorSteps +from brainles_preprocessing.modality import Modality +from brainles_preprocessing.registration import ( + ElastixRegistrator, + NiftyRegRegistrator, +) + + +class TestModalityInterpolatorGuard(unittest.TestCase): + + def setUp(self): + self.modality = Modality( + modality_name="t1", + input_path=Path("/tmp/fake_t1.nii.gz"), + raw_skull_output_path=Path("/tmp/fake_output.nii.gz"), + ) + self.fixed_image_path = Path("/tmp/fake_fixed.nii.gz") + self.registration_dir = Path("/tmp/fake_registration_dir") + + def test_register_rejects_interpolator_for_non_ants(self): + with self.assertRaises(ValueError): + self.modality.register( + registrator=NiftyRegRegistrator(), + fixed_image_path=self.fixed_image_path, + registration_dir=self.registration_dir, + moving_image_name="t1", + step=PreprocessorSteps.COREGISTERED, + interpolator="linear", + ) + + def test_transform_rejects_interpolator_for_elastix(self): + with self.assertRaises(ValueError): + self.modality.transform( + registrator=ElastixRegistrator(), + fixed_image_path=self.fixed_image_path, + registration_dir_path=self.registration_dir, + moving_image_name="t1", + transformation_matrix_path=self.registration_dir / "M_t1", + step=PreprocessorSteps.COREGISTERED, + interpolator="linear", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_registrators.py b/tests/test_registrators.py index 2f8b7b4..68e66eb 100644 --- a/tests/test_registrators.py +++ b/tests/test_registrators.py @@ -1,5 +1,6 @@ from pathlib import Path import tempfile +from unittest.mock import patch import numpy as np from registrator_base import RegistratorBase @@ -19,6 +20,33 @@ def get_registrator(self): def get_method_and_extension(self): return "ants", "mat" + def test_normalize_interpolator(self): + self.assertEqual( + ANTsRegistrator.normalize_interpolator("nearestNeighbor"), "nearestNeighbor" + ) + + @patch("brainles_preprocessing.registration.ANTs.ANTs.ants") + def test_transform_interpolator_precedence(self, mock_ants): + registrator = ANTsRegistrator(transformation_params={"interpolator": "linear"}) + kwargs = dict( + fixed_image_path=self.fixed_image, + moving_image_path=self.moving_image, + transformed_image_path=self.output_dir + "/ants_precedence.nii.gz", + matrix_path=self.transform_matrix, + log_file_path=self.output_dir + "/ants_precedence.log", + ) + + registrator.transform(**kwargs) + self.assertEqual( + mock_ants.apply_transforms.call_args.kwargs["interpolator"], "linear" + ) + + registrator.transform(interpolator="nearestNeighbor", **kwargs) + self.assertEqual( + mock_ants.apply_transforms.call_args.kwargs["interpolator"], + "nearestNeighbor", + ) + class TestNiftyRegRegistrator(RegistratorBase, unittest.TestCase): def get_registrator(self): @@ -27,6 +55,11 @@ def get_registrator(self): def get_method_and_extension(self): return "niftyreg", "txt" + def test_normalize_interpolator(self): + self.assertEqual( + NiftyRegRegistrator.normalize_interpolator("nearestNeighbor"), "0" + ) + def test_invert_affine_transform(self): with tempfile.TemporaryDirectory() as tmpdir: # Erstelle eine Testmatrix @@ -100,6 +133,10 @@ def get_registrator(self): def get_method_and_extension(self): return "elastix", "txt" + def test_normalize_interpolator_raises(self): + with self.assertRaises(ValueError): + ElastixRegistrator.normalize_interpolator("nearestNeighbor") + class TestGreedyRegistrator(RegistratorBase, unittest.TestCase): def get_registrator(self):