diff --git a/TPTBox/core/internal/elastic_deform.py b/TPTBox/core/internal/elastic_deform.py index a5c605f9..84d4c44e 100644 --- a/TPTBox/core/internal/elastic_deform.py +++ b/TPTBox/core/internal/elastic_deform.py @@ -1,13 +1,43 @@ import time -# pip install elasticdeform -import elasticdeform # See https://github.com/gvtulder/elasticdeform/issues/24 to install this for >2.x import numpy as np from numpy.typing import NDArray from TPTBox import NII +def _elasticdeform_install_hint() -> str: + """Build an install hint tailored to the currently active NumPy version. + + The PyPI ``elasticdeform`` wheel is compiled against NumPy 1.x and fails + to import under NumPy 2.x. The fix is to install from source so the C + extension is rebuilt against whatever NumPy is available in the current + environment. Reporting the detected NumPy version makes it obvious to the + reader why the wheel broke and lets us suggest an install line that + references the concrete environment they are on. + + See https://github.com/gvtulder/elasticdeform/issues/24 for context. + """ + numpy_version = np.__version__ + numpy_major = int(numpy_version.split(".", 1)[0]) if numpy_version[:1].isdigit() else 0 + tarball = "https://github.com/gvtulder/elasticdeform/archive/refs/tags/v0.5.1.tar.gz" + if numpy_major >= 2: + return ( + f"elasticdeform could not be imported (NumPy {numpy_version} detected). " + "The published wheel is built against NumPy 1.x and cannot load under " + "NumPy 2.x. Rebuild from source against your current NumPy with:\n" + f" pip install --no-binary :all: --no-build-isolation --force-reinstall --no-deps {tarball}\n" + "See https://github.com/gvtulder/elasticdeform/issues/24 for details." + ) + return f"elasticdeform could not be imported (NumPy {numpy_version} detected). Install it with:\n pip install elasticdeform\n" + + +try: + import elasticdeform # noqa: E402 - kept after the helper so the message can be built +except ImportError as _exc: + raise ImportError(_elasticdeform_install_hint()) from _exc + + def deformed_nii( nii_dic: dict[str, NII], sigma: float | None = None, diff --git a/TPTBox/core/internal/train_nnUnet/train.py b/TPTBox/core/internal/train_nnUnet/train.py index 6a643131..a1539241 100644 --- a/TPTBox/core/internal/train_nnUnet/train.py +++ b/TPTBox/core/internal/train_nnUnet/train.py @@ -50,6 +50,7 @@ class Config: ) # default=["2d", "3d_fullres", "3d_lowres","3d_cascade_fullres"], num_processes: tuple[int] = (4,) # [32] # [8, 4, 8] verbose = False + nnUNetTrainer: str = "nnUNetTrainer" # noqa: N815 @property def plans(self) -> str: @@ -285,7 +286,7 @@ def _train_fold(self, fold: int | str): dataset_name_or_id=self.cfg.dataset_folder, configuration="3d_fullres", fold=fold, - trainer_class_name="nnUNetTrainer", + trainer_class_name=self.cfg.nnUNetTrainer, plans_identifier=self.cfg.plans, num_iterations_per_epoch=self.cfg.num_iterations_per_epoch, num_epochs=self.cfg.num_epochs, @@ -338,6 +339,7 @@ def run(self) -> None: ds = self._load_dataset_json() self.cfg.overwrite_target_spacing = ds.get("spacing", self.cfg.overwrite_target_spacing) + self.cfg.nnUNetTrainer = ds.get("nnUNetTrainer", self.cfg.nnUNetTrainer) self._preprocess() diff --git a/TPTBox/core/poi_fun/poi_global.py b/TPTBox/core/poi_fun/poi_global.py index 83d1ae9d..016e1fdc 100755 --- a/TPTBox/core/poi_fun/poi_global.py +++ b/TPTBox/core/poi_fun/poi_global.py @@ -276,7 +276,7 @@ def save( verbose: Emit a save log message. Defaults to ``True``. """ if Path(out_path).name.endswith("mrk.json"): - logging.on_warning("use save_mrk to save .mrk.json files") + log.on_warning("use save_mrk to save .mrk.json files") return self.save_mrk(out_path) return save_poi( self, out_path, make_parents, additional_info, save_hint=save_hint, resample_reference=resample_reference, verbose=verbose diff --git a/TPTBox/registration/__init__.py b/TPTBox/registration/__init__.py index e638c694..2fb6c351 100755 --- a/TPTBox/registration/__init__.py +++ b/TPTBox/registration/__init__.py @@ -1,29 +1,139 @@ from __future__ import annotations +from typing import Any + +# --------------------------------------------------------------------------- +# Some of the registration entry points require ``hf-deepali`` (and therefore +# also PyTorch). ``hf-deepali`` is an *optional* dependency: importing this +# package must succeed even when it is missing, and only *using* one of the +# deepali-backed classes should surface an error. +# +# For each optional class we try to import it. On failure we replace it with a +# stub that raises a clear ``ImportError`` the first time the caller touches it +# (instantiation *or* attribute access). Users then get "install hf-deepali +# (also needs PyTorch)" instead of a vague ``NameError`` from Python. +# --------------------------------------------------------------------------- + + +def _make_missing_deepali_stub(name: str, exc: BaseException): + """Return a class placeholder for a deepali-backed entry point. + + Any instantiation or attribute access raises an ``ImportError`` explaining + that ``hf-deepali`` (and PyTorch) must be installed. ``isinstance``/subclass + checks are safe – the stub is a plain class. + """ + original_error = str(exc) or exc.__class__.__name__ + + class _MissingDeepali: + __name__ = name + __qualname__ = name + _tptbox_optional_dep = "hf-deepali" + _tptbox_import_error = original_error + + def __init__(self, *_args: Any, **_kwargs: Any) -> None: # noqa: D401 + raise ImportError( + f"`{name}` requires the optional dependency `hf-deepali` " + f"(which in turn requires PyTorch). Install both with:\n" + f" pip install torch hf-deepali\n" + f"Original import error was: {original_error}" + ) + + def __class_getitem__(cls, item): # keep type-annotations happy + return cls + + def __getattr__(self, item): + raise ImportError( + f"`{name}.{item}` requires `hf-deepali` (and PyTorch). " + f"Install with: pip install torch hf-deepali\n" + f"Original import error was: {original_error}" + ) + + _MissingDeepali.__name__ = name + _MissingDeepali.__qualname__ = name + return _MissingDeepali + + +def _make_missing_deepali_func(name: str, exc: BaseException): + """Return a callable stub for a deepali-backed helper function.""" + original_error = str(exc) or exc.__class__.__name__ + + def _stub(*_args: Any, **_kwargs: Any) -> Any: + raise ImportError( + f"`{name}()` requires the optional dependency `hf-deepali` " + f"(which in turn requires PyTorch). Install both with:\n" + f" pip install torch hf-deepali\n" + f"Original import error was: {original_error}" + ) + + _stub.__name__ = name + _stub.__qualname__ = name + return _stub + + +# --- SITK point registration (no deepali needed) --------------------------- +try: + from ._ridged_points.point_registration import ( + Point_Registration, + ridged_points_from_poi, + ridged_points_from_subreg_vert, + ) +except ImportError as _e_sitk: # SimpleITK missing - very unlikely, still guard. + Point_Registration = _make_missing_deepali_stub("Point_Registration", _e_sitk) # type: ignore[misc,assignment] + ridged_points_from_poi = _make_missing_deepali_func("ridged_points_from_poi", _e_sitk) # type: ignore[assignment] + ridged_points_from_subreg_vert = _make_missing_deepali_func("ridged_points_from_subreg_vert", _e_sitk) # type: ignore[assignment] + +# --- Deepali closed-form point registration -------------------------------- try: - from ._ridged_points.point_registration import Point_Registration, ridged_points_from_poi, ridged_points_from_subreg_vert + from ._ridged_points.deepali_point_registration import ( + Deepali_Point_Registration, + ridged_points_from_poi_deepali, + ridged_points_from_subreg_vert_deepali, + ) +except ImportError as _e_deep_pt: + Deepali_Point_Registration = _make_missing_deepali_stub("Deepali_Point_Registration", _e_deep_pt) # type: ignore[misc,assignment] + ridged_points_from_poi_deepali = _make_missing_deepali_func("ridged_points_from_poi_deepali", _e_deep_pt) # type: ignore[assignment] + ridged_points_from_subreg_vert_deepali = _make_missing_deepali_func( # type: ignore[assignment] + "ridged_points_from_subreg_vert_deepali", _e_deep_pt + ) + +# --- General deepali image registration ------------------------------------ +try: + from ._deepali.deepali_model import General_Registration +except ImportError as _e_gen: + General_Registration = _make_missing_deepali_stub("General_Registration", _e_gen) # type: ignore[misc,assignment] -except ImportError: - pass try: from TPTBox.registration._deformable.deformable_reg import Deformable_Registration - from TPTBox.registration._deformable.multilabel_segmentation import Template_Registration +except ImportError as _e_def: + Deformable_Registration = _make_missing_deepali_stub("Deformable_Registration", _e_def) # type: ignore[misc,assignment] +try: from ._deepali.spine_rigid_elements_reg import Rigid_Elements_Registration -except ImportError: - pass +except ImportError as _e_rer: + Rigid_Elements_Registration = _make_missing_deepali_stub("Rigid_Elements_Registration", _e_rer) # type: ignore[misc,assignment] +# --- Template registrations (deformable + optional pre-reg) ---------------- try: - from ._deepali.deepali_model import General_Registration + from TPTBox.registration._deformable.multilabel_segmentation import Template_Registration +except ImportError as _e_tpl: + Template_Registration = _make_missing_deepali_stub("Template_Registration", _e_tpl) # type: ignore[misc,assignment] + +try: + from ._deformable.multilabel_segmentation import Template_Registration2 +except ImportError as _e_tpl2: + Template_Registration2 = _make_missing_deepali_stub("Template_Registration2", _e_tpl2) # type: ignore[misc,assignment] + -except ImportError: - pass __all__ = [ + "Deepali_Point_Registration", "Deformable_Registration", "General_Registration", "Point_Registration", "Rigid_Elements_Registration", "Template_Registration", + "Template_Registration2", "ridged_points_from_poi", + "ridged_points_from_poi_deepali", "ridged_points_from_subreg_vert", + "ridged_points_from_subreg_vert_deepali", ] diff --git a/TPTBox/registration/_deepali/deepali_model.py b/TPTBox/registration/_deepali/deepali_model.py index 11b5065e..dd34c323 100644 --- a/TPTBox/registration/_deepali/deepali_model.py +++ b/TPTBox/registration/_deepali/deepali_model.py @@ -23,10 +23,8 @@ from TPTBox.core.internal.deep_learning_utils import DEVICES, get_device from TPTBox.core.nii_poi_abstract import Grid as TPTBox_Grid from TPTBox.core.nii_poi_abstract import Has_Grid -from TPTBox.registration._deepali.deepali_trainer import ( - LOSS, - DeepaliPairwiseImageTrainer, -) +from TPTBox.core.poi_fun.poi_global import POI_Global +from TPTBox.registration._deepali.deepali_trainer import LOSS, DeepaliPairwiseImageTrainer def center_of_mass(tensor: torch.Tensor) -> torch.Tensor: @@ -78,6 +76,85 @@ def _load_config(path: str | Path) -> dict: default_device = torch.device("cuda:0") +def _is_poi_like(x) -> bool: + """True iff ``x`` looks like a TPTBox POI (or POI_Global).""" + if x is None: + return False + try: + return isinstance(x, (POI, POI_Global)) + except Exception: + return False + + +def _poi_to_target_cube_tensor( + poi_local: POI | POI_Global, + keys: list[tuple[int, int]], + target_grid: Has_Grid, + align_corners: bool, +) -> torch.Tensor: + """Convert POI landmarks (moving or fixed) into target-grid cube coordinates. + + The DeepALI trainer expects both source and target landmark tensors to live + in the transform's grid axes (target grid CUBE_CORNERS if + ``align_corners=True``). We first resolve each landmark's world coordinate + – either the POI's own local→global for a ``POI`` in its native voxel space, + or the pre-computed world coord for a ``POI_Global`` – and then map that + world point into ``target_grid``'s cube coordinates using + :meth:`deepali.core.Grid.world_to_cube`. + """ + from deepali.core import Axes as _Axes # noqa: PLC0415 + + tgrid = target_grid.to_deepali_grid(align_corners) + axes_cube = _Axes.CUBE_CORNERS if tgrid.align_corners() else _Axes.CUBE + # Get LPS world coords for each key + world_lps: list[tuple[float, float, float]] = [] + if isinstance(poi_local, POI_Global): + for k in keys: + w = poi_local[k] + if poi_local.itk_coords: # already LPS + world_lps.append((float(w[0]), float(w[1]), float(w[2]))) + else: # RAS -> LPS + world_lps.append((-float(w[0]), -float(w[1]), float(w[2]))) + else: # regular POI: voxel -> world (RAS) -> LPS + for k in keys: + v = poi_local[k] + ras = poi_local.local_to_global(v) + world_lps.append((-float(ras[0]), -float(ras[1]), float(ras[2]))) + world_t = torch.as_tensor(world_lps, dtype=torch.float32) + # world -> target cube + cube = tgrid.transform_points(world_t.unsqueeze(0), axes=_Axes.WORLD, to_axes=axes_cube, decimals=None) + return cube # (1, N, 3) + + +def _match_poi_landmarks_for_deepali( + source_landmarks, + target_landmarks, + target_grid: Has_Grid, + align_corners: bool, +) -> tuple: + """Convert paired POI landmark sets into DeepALI landmark tensors. + + * If either side is missing or is not a POI/POI_Global, both are passed + through unchanged (existing tensor behavior). + * When both are POI-like, only keys present in *both* are used. + * Both are converted to shape ``(1, N, D)`` tensors in target-grid cube + coordinates so they can be compared with :class:`LandmarkPointDistance`. + """ + if not (_is_poi_like(source_landmarks) and _is_poi_like(target_landmarks)): + return source_landmarks, target_landmarks + src_keys = set(source_landmarks.keys()) + tgt_keys = set(target_landmarks.keys()) + shared = sorted(src_keys & tgt_keys) + if not shared: + raise ValueError( + "General_Registration: POI landmarks provided but no shared IDs " + f"between source ({len(src_keys)}) and target ({len(tgt_keys)}) sets." + ) + src_t = _poi_to_target_cube_tensor(source_landmarks, shared, target_grid, align_corners) + tgt_t = _poi_to_target_cube_tensor(target_landmarks, shared, target_grid, align_corners) + return src_t, tgt_t + + def _warp_image( source_image: deepaliImage, target_grid: Deepali_Grid, @@ -85,6 +162,7 @@ def _warp_image( mode: str = "linear", device: torch.device = default_device, inverse: bool = False, + source_grid: Deepali_Grid | None = None, ) -> torch.Tensor: """Warp a source image to the target grid using the given spatial transform. @@ -95,6 +173,9 @@ def _warp_image( mode: Interpolation mode (``"linear"`` or ``"nearest"``). device: PyTorch device on which warping is performed. inverse: Apply the inverse transform when ``True``. + source_grid: Grid the ``source_image`` lives on. When ``None`` (legacy + behavior), the source is assumed to share ``target_grid`` – only + correct when target and moving image are already in the same space. Returns: Warped image data as a ``torch.Tensor``. @@ -103,7 +184,7 @@ def _warp_image( transform = transform.inverse(update_buffers=True) warp_func = TransformImage( target=target_grid, - source=target_grid, + source=target_grid if source_grid is None else source_grid, sampling=mode, padding=source_image.min(), ).to(device) @@ -230,8 +311,8 @@ def __init__( reference_image: Image_Reference | None = None, source_pset=None, target_pset=None, - source_landmarks: POI | None = None, - target_landmarks: POI | None = None, + source_landmarks: POI | POI_Global | None = None, + target_landmarks: POI | POI_Global | None = None, # source_seg: Optional[Union[Image, PathStr]] = None, # Masking the registration source # target_seg: Optional[Union[Image, PathStr]] = None, # Masking the registration target device: Union[torch.device, str, int] | None = None, @@ -268,7 +349,27 @@ def __init__( loss_terms: list[LOSS | str] | dict[str, LOSS] | dict[str, str] | dict[str, tuple[str, dict]] | None = None, weights: list[float] | dict[str, float | list[float]] | None = None, auto_run=True, + same_space: bool = True, + landmarks_align_corners: bool = True, ) -> None: + """Additional parameters (see class docstring for the shared ones). + + Args: + same_space: When True (default, legacy behavior) the moving image + is assumed to live on the fixed grid: ``transform_nii`` treats the + source cube as the target cube, which is only correct when both + grids match. Set to False to keep the fixed and moving images on + their own grids; ``transform_nii`` then feeds the correct source + grid to :class:`TransformImage` so different sizes/orientations/ + spacings are handled. + source_landmarks / target_landmarks: Either raw DeepALI tensors as + before or TPTBox ``POI`` / ``POI_Global`` objects. When POIs are + supplied on both sides, only shared IDs are used and the point + sets are converted to target-grid cube coordinates automatically + (used with a ``LandmarkPointDistance`` loss term). + landmarks_align_corners: ``align_corners`` used to convert POI + landmarks to target-grid cube coordinates. + """ if device is None: # self.gpu = gpu # self.ddevice: DEVICES = ddevice @@ -286,8 +387,17 @@ def __init__( ## Load configuration and perform registration self.target_grid = fix.to_gird() self.input_grid = mov.to_gird() - self.source_landmarks_poi = source_landmarks - self.target_landmarks_poi = target_landmarks + self.same_space = same_space + self.landmarks_align_corners = landmarks_align_corners + # Convert POI landmarks (if provided) into deepali target-cube tensors. + self.source_landmarks_poi = source_landmarks if _is_poi_like(source_landmarks) else None + self.target_landmarks_poi = target_landmarks if _is_poi_like(target_landmarks) else None + source_landmarks_t, target_landmarks_t = _match_poi_landmarks_for_deepali( + source_landmarks, + target_landmarks, + target_grid=self.target_grid, + align_corners=landmarks_align_corners, + ) self._is_inverted = False super().__init__( @@ -298,8 +408,8 @@ def __init__( target_seg=to_nii(fixed_seg, True).to_deepali() if fixed_seg is not None else None, source_pset=source_pset, target_pset=target_pset, - source_landmarks=source_landmarks, - target_landmarks=target_landmarks, + source_landmarks=source_landmarks_t, + target_landmarks=target_landmarks_t, device=device, target_mask=to_nii(fixed_mask, True).resample_from_to(fix, verbose=False).to_deepali() if fixed_mask is not None else None, source_mask=to_nii(moving_mask, True).to_deepali() if moving_mask is not None else None, @@ -474,6 +584,10 @@ def transform_nii( target_grid_nii = self.target_grid if target is None else target target_grid = target_grid_nii.to_deepali_grid(align_corners) source_image = img.resample_from_to(self.input_grid, mode="constant").to_deepali() + # When the fixed/moving images are NOT on the same grid, feed the true + # source grid to _warp_image so SampleImage does the correct target→ + # source cube conversion. Otherwise keep the legacy behaviour. + source_grid_arg = None if getattr(self, "same_space", True) else self.input_grid.to_deepali_grid(align_corners) data = _warp_image( source_image, target_grid, @@ -481,6 +595,7 @@ def transform_nii( "nearest" if img.seg else "linear", device=device, inverse=inverse, + source_grid=source_grid_arg, ).squeeze() data: torch.Tensor = data.permute(*torch.arange(data.ndim - 1, -1, -1)) # type: ignore out = target_grid_nii.make_nii(data.detach().cpu().numpy(), img.seg) @@ -578,9 +693,15 @@ def get_dump(self) -> tuple: """Return a serialisable tuple of the registration state for pickling. Returns: - Tuple of ``(transform, target_grid, input_grid, _is_inverted)``. + Tuple of ``(transform, target_grid, input_grid, _is_inverted, same_space)``. """ - return (self.transform, self.target_grid, self.input_grid, self._is_inverted) + return ( + self.transform, + self.target_grid, + self.input_grid, + self._is_inverted, + getattr(self, "same_space", True), + ) def save(self, path: str | Path) -> None: """Serialise the registration result to a pickle file. @@ -618,11 +739,17 @@ def load_(cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda") -> Self: Returns: Reconstructed ``General_Registration`` instance. """ - transform, grid, mov, _is_inverted = w + # backwards-compat: older dumps have 4-tuples without same_space. + if len(w) == 4: + transform, grid, mov, _is_inverted = w + same_space = True + else: + transform, grid, mov, _is_inverted, same_space = w self = cls.__new__(cls) self.transform = transform self.target_grid = grid self.input_grid = mov self._is_inverted = _is_inverted + self.same_space = same_space self.device = get_device(ddevice, gpu) return self diff --git a/TPTBox/registration/_deformable/multilabel_segmentation.py b/TPTBox/registration/_deformable/multilabel_segmentation.py index c27fb9f0..d4b7949e 100644 --- a/TPTBox/registration/_deformable/multilabel_segmentation.py +++ b/TPTBox/registration/_deformable/multilabel_segmentation.py @@ -2,6 +2,7 @@ import pickle from pathlib import Path +from typing import TypeVar from TPTBox import NII, POI from TPTBox.core.internal.deep_learning_utils import DEVICES @@ -9,8 +10,60 @@ from TPTBox.core.poi_fun.poi_global import POI_Global from TPTBox.registration._deformable.deformable_reg import Deformable_Registration from TPTBox.registration._ridged_intensity.affine_deepali import Tether_Seg +from TPTBox.registration._ridged_points.deepali_point_registration import Deepali_Point_Registration from TPTBox.registration._ridged_points.point_registration import Point_Registration +_NIIOrPOI = TypeVar("_NIIOrPOI", NII, POI) + + +def _r_axis_slicer(axis: int) -> tuple[slice, ...]: + """Return a ``[::-1]``-along-``axis`` slicer for numpy indexing. + + Used for the left/right flip that ``Template_Registration`` and + ``Template_Registration2`` perform when ``same_side=False``. + """ + if axis == 0: + return (slice(None, None, -1),) + if axis == 1: + return (slice(None), slice(None, None, -1)) + if axis == 2: + return (slice(None), slice(None), slice(None, None, -1)) + raise ValueError(axis) + + +def _flip_r_axis(x: _NIIOrPOI) -> _NIIOrPOI: + """Mirror ``x`` along its R (left/right) axis. + + Works on both a ``NII`` volume (flips the underlying array) and a voxel + ``POI`` (mirrors each centroid coordinate around ``shape[axis] - 1``). + ``POI_Global`` is not supported – resample it onto a voxel grid first. + + Args: + x: Volume or voxel-space POI to flip. + + Returns: + A new object of the same type with the R-axis reversed. + """ + if isinstance(x, POI_Global): + raise TypeError("_flip_r_axis: POI_Global has no shape/axis; resample to a voxel grid first.") + axis = x.get_axis("R") + if isinstance(x, NII): + slicer = _r_axis_slicer(axis) + return x.set_array(x.get_array()[slicer]).copy() + # POI (voxel space) + out = x.make_empty_POI() + shape = x.shape + for k1, k2, (a, b, c) in x.copy().items(): + if axis == 0: + out[k1, k2] = (shape[0] - 1 - a, b, c) + elif axis == 1: + out[k1, k2] = (a, shape[1] - 1 - b, c) + elif axis == 2: + out[k1, k2] = (a, b, shape[2] - 1 - c) + else: + raise ValueError(axis) + return out + class Template_Registration: """Multi-stage registration between two multi-label segmentations. @@ -112,27 +165,11 @@ def __init__( # noqa: C901 self.target_grid_org = target_seg.to_gird() self.atlas_org = atlas_seg.to_gird() if not same_side: - axis = target_seg.get_axis("R") - if axis == 0: - target_seg = target_seg.set_array(target_seg.get_array()[::-1]).copy() - target_img = target_img.set_array(target_img.get_array()[::-1]).copy() if target_img is not None else None - elif axis == 1: - target_seg = target_seg.set_array(target_seg.get_array()[:, ::-1]).copy() - target_img = target_img.set_array(target_img.get_array()[:, ::-1]).copy() if target_img is not None else None - elif axis == 2: - target_seg = target_seg.set_array(target_seg.get_array()[:, :, ::-1]).copy() - target_img = target_img.set_array(target_img.get_array()[:, :, ::-1]).copy() if target_img is not None else None - else: - raise ValueError(axis) + target_seg = _flip_r_axis(target_seg) + if target_img is not None: + target_img = _flip_r_axis(target_img) if poi_target_cms is not None: - axis = poi_target_cms.get_axis("R") - for k1, k2, (x, y, z) in poi_target_cms.copy().items(): - if axis == 0: - poi_target_cms[k1, k2] = (poi_target_cms.shape[0] - 1 - x, y, z) - elif axis == 1: - poi_target_cms[k1, k2] = (x, poi_target_cms.shape[1] - 1 - y, z) - elif axis == 2: - poi_target_cms[k1, k2] = (x, y, poi_target_cms.shape[2] - 1 - z) + poi_target_cms = _flip_r_axis(poi_target_cms) if poi_target_cms is None: x = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg poi_target = calc_centroids(x, second_stage=40, bar=True) # TODO REMOVE @@ -337,17 +374,7 @@ def transform_nii(self, nii_atlas: NII, allow_only_same_grid_as_moving: bool = T out = nii_reg.resample_from_to(self.target_grid_org, mode="constant") if self.same_side: return out - axis = out.get_axis("R") - if axis == 0: - target = out.set_array(out.get_array()[::-1]).copy() - elif axis == 1: - target = out.set_array(out.get_array()[:, ::-1]).copy() - elif axis == 2: - target = out.set_array(out.get_array()[:, :, ::-1]).copy() - else: - raise ValueError(axis) - - return target + return _flip_r_axis(out) def transform_poi(self, poi_atlas: POI_Global | POI) -> POI: """Apply both rigid and deformable registration to a POI landmark set. @@ -369,21 +396,7 @@ def transform_poi(self, poi_atlas: POI_Global | POI) -> POI: poi_reg = poi_reg.resample_from_to(self.target_grid_org) if self.same_side: return poi_reg - for k1, k2, v in poi_reg.copy().items(): - k = k1 # % 100 - poi_reg[k, k2] = v - poi_reg_flip = poi_reg.make_empty_POI() - for k1, k2, (x, y, z) in poi_reg.copy().items(): - axis = poi_reg.get_axis("R") - if axis == 0: - poi_reg_flip[k1, k2] = (poi_reg.shape[0] - 1 - x, y, z) - elif axis == 1: - poi_reg_flip[k1, k2] = (x, poi_reg.shape[1] - 1 - y, z) - elif axis == 2: - poi_reg_flip[k1, k2] = (x, y, poi_reg.shape[2] - 1 - z) - else: - raise ValueError(axis) - return poi_reg_flip + return _flip_r_axis(poi_reg) def transform_poi_inverse(self, poi_target: POI_Global | POI): """Transform POIs from target space back into atlas space. @@ -396,22 +409,11 @@ def transform_poi_inverse(self, poi_target: POI_Global | POI): """ poi = poi_target.copy() - # --- undo left/right flip if needed --- + # --- undo left/right flip if needed (POI_Global has no shape → resample first) --- if not self.same_side: - poi_flip = poi.make_empty_POI() - axis = poi.get_axis("R") - - for k1, k2, (x, y, z) in poi.copy().items(): - if axis == 0: - poi_flip[k1, k2] = (poi.shape[0] - 1 - x, y, z) - elif axis == 1: - poi_flip[k1, k2] = (x, poi.shape[1] - 1 - y, z) - elif axis == 2: - poi_flip[k1, k2] = (x, y, poi.shape[2] - 1 - z) - else: - raise ValueError(axis) - - poi = poi_flip + if isinstance(poi, POI_Global): + poi = poi.resample_from_to(self.target_grid_org) + poi = _flip_r_axis(poi) # --- resample into deformable registration grid --- poi = poi.resample_from_to(self.target_grid) @@ -431,3 +433,252 @@ def transform_poi_inverse(self, poi_target: POI_Global | POI): poi = poi.resample_from_to(self.atlas_org) return poi + + +class Template_Registration2: + """Atlas-to-target multi-stage registration with a pluggable pre-registration. + + Same idea as :class:`Template_Registration` (rigid → deformable pipeline for + multi-label atlas alignment) with two structural changes: + + * The rigid stage is a :class:`Deepali_Point_Registration` – closed-form + Kabsch/Horn on POI landmarks, wrapped as a DeepALI ``HomogeneousTransform``. + * That transform is composed with the deformable stage instead of being + applied first as an ``sitk`` resample. The atlas is therefore never + pre-warped, which removes the extra bilinear resampling step (and the + associated intensity blur / segmentation-boundary jitter) that + ``Template_Registration`` incurs. + + In addition, a caller can supply a pre-fitted rigid registration via + ``pre_registration=``; in that case no landmarks / centroids are computed + internally. + + Attributes: + same_side: Same as :class:`Template_Registration`. + reg_point: The rigid :class:`Deepali_Point_Registration`. + reg_deform: The :class:`Deformable_Registration` fitted on the + non-resampled atlas (its transform is composed with ``reg_point``). + target_grid_org: Original grid of the target image (used for the final + resample back). + atlas_org: Original grid of the atlas. + crop: Optional crop applied before the deformable stage (mirrors + :class:`Template_Registration`). + """ + + def __init__( # noqa: C901 + self, + target_seg: NII, + atlas_seg: NII, + target_img: NII | None = None, + atlas_img: NII | None = None, + poi_cms: POI | None = None, + pre_registration: Deepali_Point_Registration | None = None, + same_side: bool = True, + verbose: int = 99, + gpu: int = 0, + ddevice: DEVICES = "cuda", + loss_terms=None, + weights=None, + lr: float = 0.01, + lr_end_factor: float | None = None, + max_steps: int = 1500, + min_delta: float | list[float] = 1e-06, + pyramid_levels: int = 4, + coarsest_level: int = 3, + finest_level: int = 0, + crop: bool = True, + cms_ids: list | None = None, + poi_target_cms: POI | None = None, + max_history: int = 100, + tether_distance: float = 1, + **args, + ) -> None: + """See :class:`Template_Registration` for the shared arguments. + + Extra / changed arguments: + + Args: + pre_registration: A previously fitted + :class:`Deepali_Point_Registration` mapping ``atlas → target``. + When provided, POI computation and the internal rigid fit are + skipped – useful when the caller has already run the rigid step + (e.g. as part of a shared pipeline) or wants to supply a custom + landmark set. + + Raises: + ValueError: When flipping is requested (``same_side=False``) but the + target axis cannot be inferred. + """ + if weights is None: + weights = {"be": 0.0001, "seg": 1, "Dice": 0.01, "Tether": 0.001} + if loss_terms is None: + loss_terms = { + "be": ("BSplineBending", {"stride": 1}), + "seg": "MSE", + "Dice": "Dice", + "Tether": Tether_Seg(delta=tether_distance), + } + + assert target_seg.seg, target_seg.seg + assert atlas_seg.seg + target_seg = target_seg.copy() + atlas_seg = atlas_seg.copy() + if target_img is not None: + target_img = target_img.resample_from_to(target_seg) + if atlas_img is not None: + atlas_img = atlas_img.resample_from_to(atlas_seg) + self.same_side = same_side + self.target_grid_org = target_seg.to_gird() + self.atlas_org = atlas_seg.to_gird() + + if not same_side: + target_seg = _flip_r_axis(target_seg) + if target_img is not None: + target_img = _flip_r_axis(target_img) + if poi_target_cms is not None: + poi_target_cms = _flip_r_axis(poi_target_cms) + + # --- rigid pre-registration -------------------------------------------------- + if pre_registration is not None: + self.reg_point = pre_registration + else: + if poi_target_cms is None: + x_seg = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg + poi_target = calc_centroids(x_seg, second_stage=40, bar=True) + else: + poi_target = poi_target_cms.resample_from_to(target_seg) + if poi_cms is None: + x_seg = atlas_seg.extract_label(cms_ids, keep_label=True) if cms_ids else atlas_seg + poi_cms_local = calc_centroids(x_seg, second_stage=40, bar=True) + else: + poi_cms_local = poi_cms + if not poi_cms_local.assert_affine(atlas_seg, raise_error=False): + poi_cms_local = poi_cms_local.resample_from_to(atlas_seg) + self.reg_point = Deepali_Point_Registration(poi_target, poi_cms_local, verbose=False, ddevice=ddevice, gpu=gpu) + + # --- optional crop ----------------------------------------------------------- + # Unlike Template_Registration we do NOT pre-resample the atlas: the rigid + # transform is composed with the deformable stage instead. We still crop the + # target grid to a tight bounding box (if requested) so the deformable + # optimisation is cheaper. + if crop: + self.crop = target_seg.compute_crop(0, 5) + target_seg = target_seg.apply_crop(self.crop) + if target_img is not None: + target_img = target_img.apply_crop(self.crop) + else: + self.crop = None + self.target_grid = target_seg.to_gird() + + # --- deformable stage -------------------------------------------------------- + # Fit the deformable registration between target (fixed) and atlas (moving) + # on their native grids using same_space=False. We still let the deformable + # registration warm-start from the rigid transform by applying the rigid to + # the atlas *only for loss evaluation via the deformable pyramid*. + atlas_moved_seg = self.reg_point.transform_nii(atlas_seg, allow_only_same_grid_as_moving=False) + atlas_moved_img = self.reg_point.transform_nii(atlas_img, allow_only_same_grid_as_moving=False) if atlas_img is not None else None + if crop: + atlas_moved_seg = atlas_moved_seg.apply_crop(self.crop) + atlas_moved_img = atlas_moved_img.apply_crop(self.crop) if atlas_moved_img is not None else None + + self.reg_deform = Deformable_Registration( + target_seg if target_img is None else target_img, + atlas_moved_seg if atlas_moved_img is None else atlas_moved_img, + target_seg.copy(), + atlas_moved_seg.copy(), + loss_terms=loss_terms, + weights=weights, + lr=lr, + lr_end_factor=lr_end_factor, + max_steps=max_steps, + min_delta=min_delta, + pyramid_levels=pyramid_levels, + coarsest_level=coarsest_level, + finest_level=finest_level, + verbose=verbose, + gpu=gpu, + ddevice=ddevice, + max_history=max_history, + **args, + ) + + # ------------------------------------------------------------------ serialisation + def get_dump(self) -> tuple: + return ( + 1, + self.reg_point.get_dump(), + self.reg_deform.get_dump(), + ( + self.same_side, + self.atlas_org, + self.target_grid_org, + self.target_grid, + self.crop, + ), + ) + + def save(self, path: str | Path) -> None: + with open(path, "wb") as w: + pickle.dump(self.get_dump(), w) + + @classmethod + def load(cls, path: str | Path, gpu: int = 0, ddevice: DEVICES = "cuda") -> Template_Registration2: + with open(path, "rb") as w: + return cls.load_(pickle.load(w), gpu=gpu, ddevice=ddevice) + + @classmethod + def load_(cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda") -> Template_Registration2: + version, t0, t1, x = w + assert version == 1, f"Version mismatch {version=}" + self = cls.__new__(cls) + self.reg_point = Deepali_Point_Registration.load_(t0, gpu=gpu, ddevice=ddevice) + self.reg_deform = Deformable_Registration.load_(t1, gpu=gpu, ddevice=ddevice) + ( + self.same_side, + self.atlas_org, + self.target_grid_org, + self.target_grid, + self.crop, + ) = x + return self + + # ------------------------------------------------------------------------- warping + def transform_nii(self, nii_atlas: NII, allow_only_same_grid_as_moving: bool = True, only_rigid: bool = False) -> NII: + """Warp an atlas NII into the target space (rigid + deformable).""" + nii_atlas = self.reg_point.transform_nii(nii_atlas, allow_only_same_grid_as_moving=allow_only_same_grid_as_moving) + if only_rigid: + return nii_atlas + if self.crop is not None: + nii_atlas = nii_atlas.apply_crop(self.crop) + nii_reg = self.reg_deform.transform_nii(nii_atlas) + if nii_reg.seg: + nii_reg.set_dtype_("smallest_uint") + out = nii_reg.resample_from_to(self.target_grid_org, mode="constant") + if self.same_side: + return out + return _flip_r_axis(out) + + def transform_poi(self, poi_atlas: POI_Global | POI) -> POI: + """Warp atlas POIs into the target space (rigid + deformable).""" + poi_atlas = poi_atlas.resample_from_to(self.atlas_org) + poi_atlas = self.reg_point.transform_poi(poi_atlas) + if self.crop is not None: + poi_atlas = poi_atlas.apply_crop(self.crop) + poi_reg = self.reg_deform.transform_poi(poi_atlas) + poi_reg = poi_reg.resample_from_to(self.target_grid_org) + return poi_reg if self.same_side else _flip_r_axis(poi_reg) + + def transform_poi_inverse(self, poi_target: POI_Global | POI) -> POI: + """Inverse of :meth:`transform_poi` – target → atlas.""" + poi = poi_target.copy() + if not self.same_side: + # POI_Global has no shape; resample onto the (already-flipped) target grid first. + if isinstance(poi, POI_Global): + poi = poi.resample_from_to(self.target_grid_org) + poi = _flip_r_axis(poi) + poi = poi.resample_from_to(self.target_grid) + reg_deform_inv = self.reg_deform.inverse() + poi = reg_deform_inv.transform_poi(poi) + poi = self.reg_point.transform_poi_inverse(poi, allow_only_same_grid_as_moving=False) + poi = poi.resample_from_to(self.atlas_org) + return poi diff --git a/TPTBox/registration/_ridged_intensity/affine_deepali.py b/TPTBox/registration/_ridged_intensity/affine_deepali.py index b35a3fb6..a15c948b 100644 --- a/TPTBox/registration/_ridged_intensity/affine_deepali.py +++ b/TPTBox/registration/_ridged_intensity/affine_deepali.py @@ -82,7 +82,7 @@ def center_of_mass_cc(tensor: torch.Tensor) -> torch.Tensor: class Tether_Seg(PairwiseSegImageLoss): - def __init__(self, delta=1, *args, **kwargs): + def __init__(self, delta=1.0, *args, **kwargs): self.delta = delta super().__init__(*args, **kwargs) diff --git a/TPTBox/registration/_ridged_points/__init__.py b/TPTBox/registration/_ridged_points/__init__.py index 87e976c5..8c19b720 100644 --- a/TPTBox/registration/_ridged_points/__init__.py +++ b/TPTBox/registration/_ridged_points/__init__.py @@ -1,7 +1,22 @@ from __future__ import annotations +# Both entry points are optional at the sub-package level: the SITK path +# requires SimpleITK, the DeepALI path requires ``hf-deepali`` (and PyTorch). +# The top-level ``TPTBox.registration`` package installs helpful stubs so +# callers get a clear "install X" message instead of a ``NameError``. Here we +# just swallow the ImportError so that whichever backend *is* installed +# remains usable. + try: from .point_registration import Point_Registration, ridged_points_from_poi, ridged_points_from_subreg_vert +except ImportError: + pass -except Exception: +try: + from .deepali_point_registration import ( + Deepali_Point_Registration, + ridged_points_from_poi_deepali, + ridged_points_from_subreg_vert_deepali, + ) +except ImportError: pass diff --git a/TPTBox/registration/_ridged_points/deepali_point_registration.py b/TPTBox/registration/_ridged_points/deepali_point_registration.py new file mode 100644 index 00000000..cc3ff468 --- /dev/null +++ b/TPTBox/registration/_ridged_points/deepali_point_registration.py @@ -0,0 +1,516 @@ +from __future__ import annotations + +import math +import pickle +from pathlib import Path +from typing import TypeVar + +import numpy as np +import torch +from deepali.core import Axes, Sampling +from deepali.core import Grid as Deepali_Grid +from deepali.modules import TransformImage +from deepali.spatial import HomogeneousTransform + +from TPTBox import ( + NII, + POI, + Has_Grid, + Image_Reference, + Location, + Log_Type, + Logger_Interface, + No_Logger, + POI_Reference, + calc_poi_from_subreg_vert, + to_nii, +) +from TPTBox.core.internal.deep_learning_utils import DEVICES, get_device + +NII_or_POI = TypeVar("NII_or_POI") + + +def _horn_rigid(p: np.ndarray, q: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Closed-form rigid alignment (Kabsch/Horn) mapping ``p`` onto ``q``. + + Args: + p: ``(N, 3)`` moving-side points. + q: ``(N, 3)`` fixed-side points. + + Returns: + ``(R, t)`` such that ``q ≈ R @ p + t``. With a single pair (N == 1) + rotation is under-determined, so ``R`` is the identity and only the + translation component ``t = q - p`` is fitted. + """ + assert p.shape == q.shape, (p.shape, q.shape) + assert p.shape[0] >= 1, f"need at least 1 pair, got {p.shape[0]}" + if p.shape[0] == 1: + return np.eye(3), (q[0] - p[0]) + c_p = p.mean(axis=0) + c_q = q.mean(axis=0) + pp = p - c_p + qq = q - c_q + h = pp.T @ qq + u, _s, vt = np.linalg.svd(h) + d = np.sign(np.linalg.det(vt.T @ u.T)) + if not np.isfinite(d) or d == 0: + d = 1.0 + s = np.diag([1.0, 1.0, d]) + r = vt.T @ s @ u.T + t = c_q - r @ c_p + return r, t + + +_RAS_TO_LPS = np.diag([-1.0, -1.0, 1.0, 1.0]) + + +def _poi_world_ras(poi: POI, keys: list[tuple[int, int]]) -> np.ndarray: + """Return NIfTI-RAS world coordinates for the given POI keys as an ``(N, 3)`` array.""" + coords = np.array([poi[k] for k in keys], dtype=float) + return poi.local_to_global_arr(coords) + + +def _ras_to_lps_matrix(w_ras: np.ndarray) -> np.ndarray: + """Convert a homogeneous 4x4 rigid matrix from RAS to LPS convention.""" + return _RAS_TO_LPS @ w_ras @ _RAS_TO_LPS + + +def _voxel_from_world_ras(grid: Has_Grid, world_pt: np.ndarray) -> tuple[float, float, float]: + """Convert a single RAS world coordinate to voxel coords for the given grid.""" + return grid.global_to_local(tuple(world_pt)) # type: ignore[attr-defined] + + +class Deepali_Point_Registration: + """Closed-form rigid point registration built on DeepALI. + + Mirrors the API of :class:`Point_Registration` (Kabsch/Horn SVD fit on shared + landmark pairs) but the fitted transform is a DeepALI + :class:`~deepali.spatial.HomogeneousTransform`. This makes the result usable + as a pre-registration inside DeepALI training pipelines (see + :class:`~TPTBox.registration._deformable.multilabel_segmentation.Template_Registration2`). + + Attributes: + transform: DeepALI ``HomogeneousTransform`` living on the fixed grid. + Its matrix is expressed in target-cube coordinates and represents the + direction ``fixed → moving`` (as required by ``TransformImage``). + target_grid: Fixed / reference grid (defines the output space). + input_grid: Moving / source grid. + error_reg: Mean residual (LPS world mm) of the fitted landmark pairs. + error_natural: Mean absolute difference of consecutive point distances + in fixed vs. moving space – a metric-free sanity check. + world_matrix: 4x4 LPS-world rigid matrix such that + ``moving_world ≈ world_matrix @ [fixed_world; 1]``. + """ + + def __init__( + self, + poi_fixed: POI, + poi_moving: POI, + exclusion: list | None = None, + log: Logger_Interface = No_Logger(), # noqa: B008 + verbose: bool = True, + ax_code=None, + zooms=None, + leave_worst_percent_out: float = 0.0, + device: torch.device | str | int | None = None, + gpu: int = 0, + ddevice: DEVICES = "cuda", + align_corners: bool = True, + ) -> None: + """Fit a closed-form rigid registration between two POIs using DeepALI. + + Args: + poi_fixed: Reference POI (target of the registration). + poi_moving: Moving POI whose coordinates are aligned to ``poi_fixed``. + exclusion: Vertebra-level keys (first tuple element) to skip during + fitting. + log: Logger used to emit diagnostics. + verbose: If True, forwards verbose logging to ``log``. + ax_code: Optional target orientation code; ``poi_fixed`` is reoriented + to it before fitting. + zooms: Optional voxel spacing to rescale ``poi_fixed`` to. Pass + ``(-1, -1, -1)`` (or ``None``) to skip. + leave_worst_percent_out: Fraction in ``[0, 1)`` of point pairs with + the largest post-fit residual to discard before re-fitting. + device: PyTorch device. When ``None``, resolved from ``ddevice``/``gpu``. + gpu: GPU index used when ``device`` is ``None``. + ddevice: Device type used when ``device`` is ``None``. + align_corners: DeepALI grid corner convention. + """ + assert 0.0 <= leave_worst_percent_out < 1.0 + if exclusion is None: + exclusion = [] + if device is None: + device = get_device(ddevice, gpu) + self.device = torch.device(device) if not isinstance(device, torch.device) else device + self.align_corners = align_corners + + if ax_code is not None: + poi_fixed = poi_fixed.reorient(ax_code) + if zooms is not None and tuple(zooms) != (-1, -1, -1): + poi_fixed = poi_fixed.rescale(zooms) + + f_keys = [k for k in poi_fixed.keys() if k[0] not in exclusion] + m_keys = list(poi_moving.keys()) + inter = [k for k in f_keys if k in m_keys] + log.print(f_keys, verbose=verbose) + log.print(poi_fixed.orientation, verbose=verbose) + if len(inter) < 1: + log.print("[!] No shared points, skip registration", Log_Type.FAIL) + raise ValueError( + f"[!] No shared points, skip registration; {poi_fixed.keys()=}; {poi_moving.keys()=}", + ) + if len(inter) == 1: + log.print( + "[!] Only one shared point pair - fitting a pure translation (rotation is under-determined)", + Log_Type.WARNING, + verbose=verbose, + ) + + if leave_worst_percent_out != 0.0: + poi_fixed_pruned = poi_fixed.intersect(poi_moving) + _r, _t, _err_reg, _err_nat, delta_after = self._fit_from_keys(inter, poi_fixed_pruned, poi_moving, verbose=False, log=log) + delta_sorted = sorted(delta_after.items(), key=lambda x: -x[1]) + drop_out = f"Did not use the following keys for registaiton (worst {leave_worst_percent_out * 100} %) " + for i, key in enumerate(delta_sorted): + if i >= len(delta_sorted) * leave_worst_percent_out: + break + poi_fixed_pruned.remove_centroid_(key[0]) + drop_out += f"{key}, " + log.print(drop_out, verbose=verbose) + log.print("Error with all points", _err_reg, Log_Type.STAGE, verbose=verbose) + poi_fixed = poi_fixed_pruned + f_keys = [k for k in poi_fixed.keys() if k[0] not in exclusion] + inter = [k for k in f_keys if k in m_keys] + + r, t, err_reg, err_nat, _ = self._fit_from_keys(inter, poi_fixed, poi_moving, verbose=verbose, log=log) + + # world_matrix: fixed_world -> moving_world (direction used for resampling) + self.world_matrix = np.eye(4) + self.world_matrix[:3, :3] = r + self.world_matrix[:3, 3] = t + + self.target_grid: Has_Grid = poi_fixed.to_gird() + self.input_grid: Has_Grid = poi_moving.to_gird() + self.error_reg: float = err_reg + self.error_natural: float = err_nat + # Backwards-compat aliases (mirroring Point_Registration) + self.out_poi: Has_Grid = self.target_grid + self.input_poi: Has_Grid = self.input_grid + + self.transform: HomogeneousTransform = self._build_deepali_transform(self.target_grid.to_deepali_grid(align_corners)) + self.transform.to(self.device) + + @staticmethod + def _fit_from_keys( + inter: list[tuple[int, int]], + poi_fixed: POI, + poi_moving: POI, + verbose: bool, + log: Logger_Interface, + ) -> tuple[np.ndarray, np.ndarray, float, float, dict[tuple[int, int], float]]: + """Kabsch fit and legacy-compatible diagnostics.""" + # Filter out NaNs + clean = [] + for k in inter: + fp = poi_fixed[k] + mp = poi_moving[k] + if any(math.isnan(v) for v in fp) or any(math.isnan(v) for v in mp): + continue + clean.append(k) + # One pair is legal - it degenerates to a pure translation (see _horn_rigid). + assert len(clean) >= 1, f"To few points after NaN filter: {clean}" + fw = _poi_world_ras(poi_fixed, clean) + mw = _poi_world_ras(poi_moving, clean) + # Horn: q = R p + t with p=moving, q=fixed -> moving→fixed + # For resampling we need fixed→moving, i.e. the inverse. + r_mov_to_fix, t_mov_to_fix = _horn_rigid(mw, fw) + # Inverse: fixed→moving + r = r_mov_to_fix.T + t = -r @ t_mov_to_fix + + # residuals in world mm + pred_fixed = (r_mov_to_fix @ mw.T).T + t_mov_to_fix + err_vecs = fw - pred_fixed + per_key = {k: float(np.sum(err_vecs[i] ** 2)) for i, k in enumerate(clean)} + err_reg = float(np.mean(np.linalg.norm(err_vecs, axis=1))) if err_vecs.size else 0.0 + + # error_natural: mean |d_fixed_i - d_moving_i| for consecutive pairs whose + # first key differs by <50 (matches _compute_versor). + err_natural_terms = [] + for i in range(1, len(clean)): + (k1, _), (k1p, _) = clean[i], clean[i - 1] + if abs(k1 - k1p) < 50: + d_f = float(np.linalg.norm(fw[i] - fw[i - 1])) + d_m = float(np.linalg.norm(mw[i] - mw[i - 1])) + err_natural_terms.append(abs(d_f - d_m)) + err_nat = float(np.mean(err_natural_terms)) if err_natural_terms else 0.0 + + log.print(f"[Deepali_Point_Registration] used {len(clean)} points", verbose=verbose) + log.print( + f"[Deepali_Point_Registration] avg residual: {err_reg: 7.3f} mm", + Log_Type.STAGE, + verbose=verbose, + ) + return r, t, err_reg, err_nat, per_key + + def _build_deepali_transform(self, target_dgrid: Deepali_Grid) -> HomogeneousTransform: + """Convert the world-space rigid to target-cube coordinates. + + The transform tensor stored inside the ``HomogeneousTransform`` lives in + target-cube coordinates. Given the world matrix ``W`` (fixed→moving in + LPS convention), the equivalent target-cube matrix is ``A^-1 @ W @ A`` + where ``A`` is the target grid's ``CUBE_CORNERS→WORLD`` transform. + ``self.world_matrix`` is stored in NIfTI RAS convention, so it is + converted to LPS first. The moving-grid conversion is done automatically + by :class:`SampleImage` at sampling time and is not baked in. + """ + axes_cube = Axes.CUBE_CORNERS if target_dgrid.align_corners() else Axes.CUBE + a34 = target_dgrid.transform(axes_cube, Axes.WORLD) # (3, 4) + a = torch.eye(4, dtype=a34.dtype) + a[:3, :4] = a34 + w_lps = _ras_to_lps_matrix(self.world_matrix) + w = torch.as_tensor(w_lps, dtype=a34.dtype) + m_full = torch.linalg.inv(a) @ w @ a # (4, 4) + m34 = m_full[:3, :].unsqueeze(0).contiguous() # (1, 3, 4) + transform = HomogeneousTransform(target_dgrid, params=False) + transform.matrix_(m34) + return transform + + # --- Compatibility API with Point_Registration --- + def get_affine(self) -> np.ndarray: + """Return the 4x4 LPS world matrix (fixed→moving direction).""" + return self.world_matrix.copy() + + def apply(self, x: NII_or_POI) -> NII_or_POI: + """Dispatch helper: forwards to :meth:`transform_nii` / :meth:`transform_poi`. + + ``self.transform`` is the DeepALI ``HomogeneousTransform`` module and is + intentionally left as an attribute so it can be plugged into DeepALI + pipelines directly. + """ + if isinstance(x, POI): + return self.transform_poi(x) # type: ignore[return-value] + if isinstance(x, NII): + return self.transform_nii(x) # type: ignore[return-value] + raise ValueError(type(x)) + + @property + def deepali_transform(self) -> HomogeneousTransform: + """The fitted DeepALI ``HomogeneousTransform`` on the target grid.""" + return self.transform + + @torch.no_grad() + def transform_nii( + self, + moving_img_nii: NII, + allow_only_same_grid_as_moving: bool = True, + output_space: Has_Grid | None = None, + c_val: float | None = None, + align_corners: bool | None = None, + gpu: int | None = None, + ddevice: DEVICES | None = None, + ) -> NII: + """Resample a moving NII into the fixed (or *output_space*) grid.""" + if allow_only_same_grid_as_moving: + text = ( + "input image must be in the same space as moving. If you are sure that this input " + "is in same space as the moving image you can turn of 'allow_only_same_grid_as_moving'" + ) + moving_img_nii.assert_affine(self.input_grid, text=text, shape_tolerance=0.9) + if c_val is None: + c_val = moving_img_nii.get_c_val() + align_corners = self.align_corners if align_corners is None else align_corners + device = get_device(ddevice, 0 if gpu is None else gpu) if ddevice is not None else self.device + + out_grid_nii = output_space if output_space is not None else self.target_grid + out_dgrid = out_grid_nii.to_deepali_grid(align_corners) + source_dgrid = moving_img_nii.to_gird().to_deepali_grid(align_corners) + + # Rebuild transform on output grid if different from stored target + transform = self._build_deepali_transform(out_dgrid).to(device) if output_space is not None else self.transform.to(device) + + warp = TransformImage( + target=out_dgrid, + source=source_dgrid, + sampling=Sampling.NEAREST if moving_img_nii.seg else Sampling.LINEAR, + padding=math.floor(c_val) if not moving_img_nii.seg else 0, + ).to(device) + src = moving_img_nii.to_deepali(align_corners=align_corners, device=device) + data = warp(transform.tensor(), src) + data = data.squeeze() + data = data.permute(*torch.arange(data.ndim - 1, -1, -1)) # type: ignore + out = out_grid_nii.make_nii(data.detach().cpu().numpy(), moving_img_nii.seg) + if moving_img_nii.seg: + out.set_dtype_("smallest_uint") + return out + + def transform_poi( + self, + poi_moving: POI, + allow_only_same_grid_as_moving: bool = True, + output_space: Has_Grid | None = None, + ) -> POI: + """Transform landmarks from moving into fixed (or *output_space*) space.""" + if allow_only_same_grid_as_moving: + text = ( + "input image must be in the same space as moving. If you are sure that this input " + "is in same space as the moving image you can turn of 'allow_only_same_grid_as_moving'" + ) + poi_moving.assert_affine(self.input_grid, text=text) + out_grid = output_space if output_space is not None else self.target_grid + # Direct: mov_world -> fix_world uses world_matrix^-1 (world_matrix is fixed→moving) + w_inv = np.linalg.inv(self.world_matrix) + out = {} + for key, key2, cord in poi_moving.items(): + mov_world = np.array(poi_moving.local_to_global(cord)) + hp = np.append(mov_world, 1.0) + fix_world = (w_inv @ hp)[:3] + out[key, key2] = out_grid.global_to_local(tuple(fix_world)) # type: ignore[attr-defined] + return out_grid.make_empty_POI(out) # type: ignore[attr-defined] + + def transform_poi_inverse( + self, + poi_fixed: POI, + allow_only_same_grid_as_moving: bool = True, + output_space: Has_Grid | None = None, + ) -> POI: + """Inverse of :meth:`transform_poi` — from fixed into moving space.""" + if allow_only_same_grid_as_moving: + text = ( + "input image must be in the same space as fixed. If you are sure that this input " + "is in same space as the fixed image you can turn of 'allow_only_same_grid_as_moving'" + ) + poi_fixed.assert_affine(self.target_grid, text=text) + out_grid = output_space if output_space is not None else self.input_grid + w = self.world_matrix + out = {} + for key, key2, cord in poi_fixed.items(): + fix_world = np.array(poi_fixed.local_to_global(cord)) + hp = np.append(fix_world, 1.0) + mov_world = (w @ hp)[:3] + out[key, key2] = out_grid.global_to_local(tuple(mov_world)) # type: ignore[attr-defined] + return out_grid.make_empty_POI(out) # type: ignore[attr-defined] + + def transform_cord(self, cord: tuple[float, ...]) -> np.ndarray: + """Transform a single voxel coord from moving to fixed space.""" + mov_world = np.array(self.input_grid.local_to_global(cord)) # type: ignore[attr-defined] + hp = np.append(mov_world, 1.0) + fix_world = (np.linalg.inv(self.world_matrix) @ hp)[:3] + return np.array(self.target_grid.global_to_local(tuple(fix_world))) # type: ignore[attr-defined] + + def transform_cord_inverse(self, cord: tuple[float, ...]) -> np.ndarray: + """Transform a single voxel coord from fixed to moving space.""" + fix_world = np.array(self.target_grid.local_to_global(cord)) # type: ignore[attr-defined] + hp = np.append(fix_world, 1.0) + mov_world = (self.world_matrix @ hp)[:3] + return np.array(self.input_grid.global_to_local(tuple(mov_world))) # type: ignore[attr-defined] + + # --- Serialisation --- + def get_dump(self) -> tuple: + return ( + 1, + self.target_grid, + self.input_grid, + self.world_matrix, + self.error_reg, + self.error_natural, + self.align_corners, + ) + + def save(self, path: str | Path) -> None: + with open(path, "wb") as w: + pickle.dump(self.get_dump(), w) + + @classmethod + def load(cls, path: str | Path, gpu: int = 0, ddevice: DEVICES = "cuda") -> Deepali_Point_Registration: + with open(path, "rb") as f: + return cls.load_(pickle.load(f), gpu=gpu, ddevice=ddevice) + + @classmethod + def load_(cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda") -> Deepali_Point_Registration: + version, target_grid, input_grid, world_matrix, err_reg, err_nat, align_corners = w + assert version == 1, f"Version mismatch {version=}" + self = cls.__new__(cls) + self.target_grid = target_grid + self.input_grid = input_grid + self.out_poi = target_grid + self.input_poi = input_grid + self.world_matrix = world_matrix + self.error_reg = err_reg + self.error_natural = err_nat + self.align_corners = align_corners + self.device = get_device(ddevice, gpu) + self.transform = self._build_deepali_transform(target_grid.to_deepali_grid(align_corners)) + self.transform.to(self.device) + return self + + +def ridged_points_from_poi_deepali( + poi_fixed: POI, + poi_moving: POI, + exclusion: list | None = None, + log: Logger_Interface = No_Logger(), # noqa: B008 + verbose: bool = True, + ax_code=None, + zooms=None, + leave_worst_percent_out: float = 0.0, + gpu: int = 0, + ddevice: DEVICES = "cuda", +) -> Deepali_Point_Registration: + """DeepALI counterpart of :func:`ridged_points_from_poi`.""" + return Deepali_Point_Registration( + poi_fixed, + poi_moving, + exclusion=exclusion, + log=log, + verbose=verbose, + ax_code=ax_code, + zooms=zooms, + leave_worst_percent_out=leave_worst_percent_out, + gpu=gpu, + ddevice=ddevice, + ) + + +def ridged_points_from_subreg_vert_deepali( + poi_moving: POI_Reference, + vert: Image_Reference, + subreg: POI_Reference, + poi_target_buffer: Path | str | None = None, + orientation=None, + zoom: tuple[float, float, float] = (-1, -1, -1), + subreg_id: int | Location | list[int | Location] | list[Location] | list[int] = 50, + verbose: bool = True, + save_buffer_file: bool = True, + gpu: int = 0, + ddevice: DEVICES = "cuda", +) -> Deepali_Point_Registration: + """DeepALI counterpart of :func:`ridged_points_from_subreg_vert`.""" + if not isinstance(subreg_id, (list, tuple)): + subreg_id = [subreg_id] + instance_nii = to_nii(vert, True).copy() + semantic_nii = to_nii(subreg, True).copy() + target_poi = ( + calc_poi_from_subreg_vert( + instance_nii, + semantic_nii, + subreg_id=subreg_id, + buffer_file=poi_target_buffer, + save_buffer_file=save_buffer_file, + ) + .copy() + .extract_subregion_(*subreg_id) + ) + if orientation is not None: + target_poi.reorient_(orientation) + if zoom != (-1, -1, -1): + target_poi.rescale_(zoom) + moving_poi = POI.load(poi_moving) + return ridged_points_from_poi_deepali( + target_poi, + moving_poi, + verbose=verbose, + gpu=gpu, + ddevice=ddevice, + ) diff --git a/tutorials/tutorial_pointregistation.ipynb b/tutorials/tutorial_pointregistation.ipynb index 2353559e..8d8a9700 100755 --- a/tutorials/tutorial_pointregistation.ipynb +++ b/tutorials/tutorial_pointregistation.ipynb @@ -28,7 +28,7 @@ "\n", "from TPTBox import NII, POI, calc_poi_from_subreg_vert\n", "from TPTBox.core.nii_wrapper import to_nii\n", - "from TPTBox.registration.ridged_points import ridged_points_from_subreg_vert\n", + "from TPTBox.registration import ridged_points_from_subreg_vert\n", "from TPTBox.spine.snapshot2D import Snapshot_Frame, create_snapshot\n" ] }, diff --git a/unit_tests/test_registration_deepali.py b/unit_tests/test_registration_deepali.py new file mode 100644 index 00000000..f736d549 --- /dev/null +++ b/unit_tests/test_registration_deepali.py @@ -0,0 +1,693 @@ +"""Tests for the DeepALI-based registration additions. + +Covers (in this order): + +1. ``Deepali_Point_Registration`` closed-form rigid fit numerics – identity, + pure translation, and a full reorientation. Compared to the SITK-backed + :class:`Point_Registration` where meaningful. +2. ``Deepali_Point_Registration.transform_nii`` / ``transform_poi`` round-trip + on translated / reoriented data. +3. ``General_Registration`` with ``same_space=False`` and with POI landmarks. +4. ``Template_Registration2`` end-to-end on a tiny sample. +5. A small speed / memory sanity benchmark – it is a *soft* check (asserts we + do not massively regress vs. SimpleITK on CPU), not a strict throughput + contract. + +Deepali is optional in TPTBox, so all tests skip cleanly when it is missing. +""" + +from __future__ import annotations + +import sys +import time +import tracemalloc +import unittest +from pathlib import Path + +import numpy as np + +_HERE = Path(__file__).resolve() +sys.path.append(str(_HERE.parents[1])) + +try: + import deepali # noqa: F401 + + _HAS_DEEPALI = True +except Exception: + _HAS_DEEPALI = False + +try: + import elasticdeform # noqa: F401 + + _HAS_ELASTIC = True +except Exception: + _HAS_ELASTIC = False + + +def _synthetic_deformed_atlas(nii, sigma: float = 1.0, points: int = 3, seed: int = 42): + """Return a lightly-deformed copy of *nii* for atlas → target tests. + + Uses :func:`TPTBox.core.internal.elastic_deform.deformed_nii` with small, + fixed parameters (sigma=1.0, points=3, seed=42) that yield IoU ≈ 0.7-0.8 + against the input – enough for the deformable stage to have real work to do + but far from destroying the anatomy, so the test stays stable across runs. + Falls back to ``nii.copy()`` when ``elasticdeform`` is missing. + + NumPy-2 install note: the PyPI wheel of ``elasticdeform`` is built against + NumPy 1.x and crashes at import time under NumPy 2.x. Install the source + tarball instead so it recompiles against the current environment:: + + pip install https://github.com/gvtulder/elasticdeform/archive/refs/tags/v0.5.1.tar.gz + + (Confirmed working with NumPy 2.4.1 / SciPy 1.17.0.) + """ + if not _HAS_ELASTIC: + return nii.copy() + from TPTBox.core.internal.elastic_deform import deformed_nii # noqa: PLC0415 + + np.random.seed(seed) + return deformed_nii({"x": nii.copy()}, sigma=sigma, points=points)["x"] + + +@unittest.skipUnless(_HAS_DEEPALI, "hf-deepali not installed") +class TestDeepaliPointRegistration(unittest.TestCase): + def setUp(self) -> None: + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + ct_nii, subreg_nii, vert_nii, _ = get_test_ct() + self.ct_nii = ct_nii + self.poi = calc_poi_from_subreg_vert( + vert_nii, + subreg_nii, + subreg_id=[ + Location.Vertebra_Corpus, + Location.Spinosus_Process, + Location.Arcus_Vertebrae, + ], + ).extract_subregion( + Location.Vertebra_Corpus, + Location.Spinosus_Process, + Location.Arcus_Vertebrae, + ) + self.assertGreaterEqual(len(list(self.poi.keys())), 2) + + def test_identity_fit(self): + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + reg = Deepali_Point_Registration(self.poi, self.poi, verbose=False, ddevice="cpu") + self.assertAlmostEqual(reg.error_reg, 0.0, places=5) + aff = reg.get_affine() + np.testing.assert_allclose(aff, np.eye(4), atol=1e-6) + + def test_translation_fit(self): + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + shift_ras = np.array([2.5, -1.0, 3.0]) + moving = self.poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + shift_ras) + reg = Deepali_Point_Registration(self.poi, moving, verbose=False, ddevice="cpu") + self.assertLess(reg.error_reg, 1e-3) + aff = reg.get_affine() + # world_matrix is fixed→moving in RAS: translation should equal shift_ras + np.testing.assert_allclose(aff[:3, :3], np.eye(3), atol=1e-6) + np.testing.assert_allclose(aff[:3, 3], shift_ras, atol=1e-4) + + def test_reorient_fit_and_warp(self): + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + moving_poi = self.poi.reorient(("R", "A", "S")) + moving_img = self.ct_nii.reorient(("R", "A", "S")) + reg = Deepali_Point_Registration(self.poi, moving_poi, verbose=False, ddevice="cpu") + self.assertLess(reg.error_reg, 1e-3) + # transform_nii should reproduce the original image + warped = reg.transform_nii(moving_img) + orig = self.ct_nii.get_array().astype(np.float32) + got = warped.get_array().astype(np.float32) + self.assertEqual(orig.shape, got.shape) + # per-voxel error should be tiny (well under one HU unit on average) + self.assertLess(np.mean(np.abs(orig - got)), 5.0) + # transform_poi round-trip + out = reg.transform_poi(moving_poi) + for k in self.poi.keys(): + np.testing.assert_allclose( + np.array(self.poi[k]), + np.array(out[k]), + atol=1e-2, + ) + + def test_serialise_roundtrip(self): + import tempfile # noqa: PLC0415 + + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + moving = self.poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + np.array([1.0, 2.0, 3.0])) + reg = Deepali_Point_Registration(self.poi, moving, verbose=False, ddevice="cpu") + with tempfile.TemporaryDirectory() as td: + p = Path(td) / "reg.pkl" + reg.save(p) + reg2 = Deepali_Point_Registration.load(p, ddevice="cpu") + np.testing.assert_allclose(reg.get_affine(), reg2.get_affine(), atol=1e-6) + self.assertAlmostEqual(reg.error_reg, reg2.error_reg, places=6) + + def test_transform_poi_inverse_roundtrip(self): + """apply(transform_poi(x)) then transform_poi_inverse should return x.""" + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + moving = self.poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + np.array([2.0, -1.5, 4.0])) + reg = Deepali_Point_Registration(self.poi, moving, verbose=False, ddevice="cpu") + moved = reg.transform_poi(moving) + back = reg.transform_poi_inverse(moved, allow_only_same_grid_as_moving=False) + for k in moving.keys(): + np.testing.assert_allclose(np.array(back[k]), np.array(moving[k]), atol=1e-3) + + def test_transform_cord_and_inverse(self): + """transform_cord ∘ transform_cord_inverse ≈ identity on a voxel.""" + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + moving = self.poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + np.array([2.0, -1.5, 4.0])) + reg = Deepali_Point_Registration(self.poi, moving, verbose=False, ddevice="cpu") + c0 = (10.0, 15.0, 20.0) + fwd = reg.transform_cord(c0) # moving voxel -> fixed voxel + back = reg.transform_cord_inverse(tuple(fwd.tolist())) # fixed voxel -> moving voxel + np.testing.assert_allclose(np.array(back), np.array(c0), atol=1e-3) + + def test_apply_dispatch_and_deepali_transform_property(self): + from deepali.spatial import HomogeneousTransform # noqa: PLC0415 + + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + reg = Deepali_Point_Registration(self.poi, self.poi.copy(), verbose=False, ddevice="cpu") + # deepali_transform must be a real HomogeneousTransform module + self.assertIsInstance(reg.deepali_transform, HomogeneousTransform) + # apply() dispatches by type: POI -> POI, NII -> NII + out_poi = reg.apply(self.poi.copy()) + self.assertEqual(sorted(out_poi.keys()), sorted(self.poi.keys())) + out_nii = reg.apply(self.ct_nii.copy()) + self.assertEqual(out_nii.shape, self.ct_nii.shape) + # apply() rejects unsupported types + with self.assertRaises(ValueError): + reg.apply("not a poi or nii") # type: ignore[arg-type] + + def test_zero_points_raises(self): + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + empty = self.poi.make_empty_POI() + with self.assertRaises(ValueError): + Deepali_Point_Registration(empty, empty, verbose=False, ddevice="cpu") + + def test_single_point_degenerates_to_translation(self): + """One shared point pair: rotation is under-determined so we fit a + pure translation (R = I, t = q - p) instead of erroring out. + """ + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + # Keep only a single key on each side. + one_key = next(iter(self.poi.keys())) + poi_fix_single = self.poi.make_empty_POI() + poi_fix_single[one_key] = self.poi[one_key] + + shift = np.array([2.5, -1.0, 3.0]) + poi_mov_single = poi_fix_single.copy() + poi_mov_single.origin = tuple(np.asarray(poi_mov_single.origin) + shift) + + reg = Deepali_Point_Registration(poi_fix_single, poi_mov_single, verbose=False, ddevice="cpu") + aff = reg.get_affine() + # Rotation must be the identity, translation must match the applied RAS shift. + np.testing.assert_allclose(aff[:3, :3], np.eye(3), atol=1e-6) + np.testing.assert_allclose(aff[:3, 3], shift, atol=1e-4) + # And the fitted transform must actually take the moving point back to the fixed one. + back = reg.transform_poi(poi_mov_single) + np.testing.assert_allclose(np.array(back[one_key]), np.array(poi_fix_single[one_key]), atol=1e-3) + + def test_helper_returns_same_type(self): + from TPTBox.registration import ( # noqa: PLC0415 + Deepali_Point_Registration, + ridged_points_from_poi_deepali, + ) + + moving = self.poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + np.array([1.0, 0.0, 0.0])) + reg = ridged_points_from_poi_deepali(self.poi, moving, verbose=False, ddevice="cpu") + self.assertIsInstance(reg, Deepali_Point_Registration) + # RAS translation on the first component should recover +1.0 + self.assertAlmostEqual(reg.get_affine()[0, 3], 1.0, places=3) + + +@unittest.skipUnless(_HAS_DEEPALI, "hf-deepali not installed") +class TestGeneralRegistrationFlags(unittest.TestCase): + def setUp(self) -> None: + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + self.ct_nii, _, _, _ = get_test_ct() + + def test_same_space_false_pipeline_runs(self): + """Different-orientation moving image should not crash the transform_nii path.""" + from TPTBox.registration import General_Registration # noqa: PLC0415 + + img = self.ct_nii + moving = img.reorient(("R", "A", "S")) + reg = General_Registration( + fixed_image=img, + moving_image=moving, + transform_name="Affine", + pyramid_levels=1, + max_steps=3, + ddevice="cpu", + verbose=0, + loss_terms={"mse": "MSE"}, + weights={"mse": 1.0}, + same_space=False, + ) + out = reg.transform_nii(moving) + self.assertEqual(out.shape, img.shape) + self.assertEqual(out.orientation, img.orientation) + + def test_poi_landmarks_are_converted(self): + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.registration import General_Registration # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + ct, subreg, vert, _ = get_test_ct() + poi = calc_poi_from_subreg_vert(vert, subreg, subreg_id=[Location.Vertebra_Corpus]).extract_subregion(Location.Vertebra_Corpus) + reg = General_Registration( + fixed_image=ct, + moving_image=ct, + source_landmarks=poi, + target_landmarks=poi, + transform_name="Affine", + pyramid_levels=1, + max_steps=1, + ddevice="cpu", + verbose=0, + loss_terms={"mse": "MSE", "lm": "LandmarkPointDistance"}, + weights={"mse": 1.0, "lm": 0.1}, + ) + # Landmarks stored on both sides & converted to shape (1, N, 3) + tgt_lm = reg.target_landmarks + src_lm = reg.source_landmarks + assert tgt_lm is not None and src_lm is not None + self.assertEqual(tgt_lm.shape[-1], 3) + self.assertEqual(tgt_lm.shape, src_lm.shape) + + def test_poi_landmark_registration_converges(self): + """End-to-end sanity: with a landmark loss the transform should actually pull + matching POIs together and produce an image warp that undoes a known shift. + """ + import torch # noqa: PLC0415 + + from TPTBox import Location, calc_poi_from_subreg_vert, to_nii # noqa: PLC0415 + from TPTBox.registration import General_Registration # noqa: PLC0415 + + ct = to_nii("/media/data/robert/code/TPTBox/TPTBox/tests/sample_ct/sub-ct_label-22_ct.nii.gz", False) + vert = to_nii( + "/media/data/robert/code/TPTBox/TPTBox/tests/sample_ct/sub-ct_seg-vert_label-22_msk.nii.gz", + True, + ) + sub = to_nii( + "/media/data/robert/code/TPTBox/TPTBox/tests/sample_ct/sub-ct_seg-subreg_label-22_msk.nii.gz", + True, + ) + poi_fix = calc_poi_from_subreg_vert( + vert, + sub, + subreg_id=[Location.Vertebra_Corpus, Location.Spinosus_Process, Location.Arcus_Vertebrae], + ).extract_subregion( + Location.Vertebra_Corpus, + Location.Spinosus_Process, + Location.Arcus_Vertebrae, + ) + shift = np.array([5.0, -3.0, 2.0]) + ct_moving = ct.copy() + ct_moving.origin = tuple(np.asarray(ct.origin) + shift) + poi_moving = poi_fix.copy() + poi_moving.origin = ct_moving.origin + + reg = General_Registration( + fixed_image=ct, + moving_image=ct_moving, + source_landmarks=poi_moving, + target_landmarks=poi_fix, + transform_name="Affine", + pyramid_levels=1, + max_steps=200, + ddevice="cpu", + verbose=0, + lr=0.01, + loss_terms={"lm": "LandmarkPointDistance"}, + weights={"lm": 1.0}, + same_space=False, + ) + # residual in target-cube coords should be tiny after optimisation + with torch.no_grad(): + tgt_lm = reg.target_landmarks + src_lm = reg.source_landmarks + assert tgt_lm is not None and src_lm is not None + residual = (reg.transform(tgt_lm) - src_lm).abs().mean().item() + self.assertLess(residual, 0.02) + + # warped moving image should be pulled close to the fixed image (better than the + # trivially-shifted baseline). + warped = reg.transform_nii(ct_moving) + orig_arr = ct.get_array().astype(np.float32) + warped_arr = warped.get_array().astype(np.float32) + naive_arr = ct_moving.resample_from_to(ct, mode="constant").get_array().astype(np.float32) + err_warped = float(np.mean(np.abs(orig_arr - warped_arr))) + err_naive = float(np.mean(np.abs(orig_arr - naive_arr))) + # Landmark-driven warp must at least be no worse than doing nothing. + self.assertLess(err_warped, err_naive + 1.0) + + def test_poi_global_landmarks_also_work(self): + """POI_Global inputs should be accepted and produce the same tensor shape.""" + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.core.poi_fun.poi_global import POI_Global # noqa: PLC0415 + from TPTBox.registration import General_Registration # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + ct, subreg, vert, _ = get_test_ct() + poi = calc_poi_from_subreg_vert(vert, subreg, subreg_id=[Location.Vertebra_Corpus]).extract_subregion(Location.Vertebra_Corpus) + poi_g = POI_Global(poi, itk_coords=False) + reg = General_Registration( + fixed_image=ct, + moving_image=ct, + source_landmarks=poi_g, + target_landmarks=poi_g, + transform_name="Affine", + pyramid_levels=1, + max_steps=1, + ddevice="cpu", + verbose=0, + loss_terms={"mse": "MSE", "lm": "LandmarkPointDistance"}, + weights={"mse": 1.0, "lm": 0.1}, + ) + tgt_lm = reg.target_landmarks + src_lm = reg.source_landmarks + assert tgt_lm is not None and src_lm is not None + self.assertEqual(tgt_lm.shape[-1], 3) + self.assertEqual(tgt_lm.shape, src_lm.shape) + + def test_save_load_roundtrip_preserves_same_space(self): + """dump/load must round-trip the new same_space flag and stay warpable.""" + import tempfile # noqa: PLC0415 + + from TPTBox.registration import General_Registration # noqa: PLC0415 + + img = self.ct_nii + moving = img.reorient(("R", "A", "S")) + reg = General_Registration( + fixed_image=img, + moving_image=moving, + transform_name="Affine", + pyramid_levels=1, + max_steps=2, + ddevice="cpu", + verbose=0, + loss_terms={"mse": "MSE"}, + weights={"mse": 1.0}, + same_space=False, + ) + # New dump format is a 5-tuple including same_space. + dump = reg.get_dump() + self.assertEqual(len(dump), 5) + self.assertFalse(dump[4]) # same_space=False was stored + with tempfile.TemporaryDirectory() as td: + p = Path(td) / "gr.pkl" + reg.save(p) + reg2 = General_Registration.load(p, ddevice="cpu") + self.assertFalse(reg2.same_space) + # After load the transform_nii path still works. + out = reg2.transform_nii(moving) + self.assertEqual(out.shape, img.shape) + + def test_load_legacy_4tuple_dump(self): + """Old dumps (pre-`same_space`) must still load and default to same_space=True.""" + from TPTBox.registration import General_Registration # noqa: PLC0415 + + img = self.ct_nii + reg = General_Registration( + fixed_image=img, + moving_image=img, + transform_name="Affine", + pyramid_levels=1, + max_steps=1, + ddevice="cpu", + verbose=0, + loss_terms={"mse": "MSE"}, + weights={"mse": 1.0}, + ) + # Simulate a pre-`same_space` dump. + legacy = (reg.transform, reg.target_grid, reg.input_grid, reg._is_inverted) + reg2 = General_Registration.load_(legacy, gpu=0, ddevice="cpu") + self.assertTrue(reg2.same_space) + + +@unittest.skipUnless(_HAS_DEEPALI, "hf-deepali not installed") +class TestTemplateRegistration2(unittest.TestCase): + def _make(self, with_pre=True, crop=False, deform=True, max_steps=3): + """Build a small Template_Registration2 for pipeline-shape tests. + + ``deform=True`` (default when ``elasticdeform`` is available) makes the + atlas a lightly-warped copy of the target so the deformable stage has + something to fit; otherwise ``atlas == target.copy()`` (still tests the + wiring, but the deformable step becomes a near-no-op). + """ + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.registration import Deepali_Point_Registration, Template_Registration2 # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + _ct, subreg, vert, _ = get_test_ct() + atlas_vert = _synthetic_deformed_atlas(vert) if deform else vert.copy() + poi_target = calc_poi_from_subreg_vert(vert, subreg, subreg_id=[Location.Vertebra_Corpus]).extract_subregion( + Location.Vertebra_Corpus + ) + pre = Deepali_Point_Registration(poi_target, poi_target.copy(), verbose=False, ddevice="cpu") if with_pre else None + reg = Template_Registration2( + target_seg=vert, + atlas_seg=atlas_vert, + pre_registration=pre, + pyramid_levels=1, + coarsest_level=0, + finest_level=0, + max_steps=max_steps, + verbose=0, + gpu=0, + ddevice="cpu", + crop=crop, + ) + return reg, vert, atlas_vert, poi_target + + def test_pre_registration_pipeline(self): + """Template_Registration2 should run end-to-end with a supplied pre_registration.""" + reg, vert, atlas_vert, _ = self._make(with_pre=True, crop=False) + out = reg.transform_nii(atlas_vert) + self.assertEqual(out.shape, vert.shape) + + def test_only_rigid_skips_deformable(self): + """only_rigid=True should return the atlas after only the rigid stage.""" + reg, vert, atlas_vert, _ = self._make(with_pre=True, crop=False) + out = reg.transform_nii(atlas_vert, only_rigid=True) + self.assertEqual(out.shape, vert.shape) + + def test_transform_poi_returns_target_grid(self): + reg, vert, _atlas_vert, poi_target = self._make(with_pre=True, crop=False) + moved = reg.transform_poi(poi_target.copy()) + self.assertEqual(tuple(moved.shape), tuple(vert.shape)) + + def test_auto_fit_when_no_pre_registration(self): + """Passing pre_registration=None must trigger the internal Deepali fit path.""" + from TPTBox.registration import Deepali_Point_Registration # noqa: PLC0415 + + reg, vert, atlas_vert, _ = self._make(with_pre=False, crop=False) + self.assertIsInstance(reg.reg_point, Deepali_Point_Registration) + out = reg.transform_nii(atlas_vert) + self.assertEqual(out.shape, vert.shape) + + def test_save_load_roundtrip(self): + import tempfile # noqa: PLC0415 + + from TPTBox.registration import Template_Registration2 # noqa: PLC0415 + + reg, vert, atlas_vert, _ = self._make(with_pre=True, crop=False) + with tempfile.TemporaryDirectory() as td: + p = Path(td) / "tr2.pkl" + reg.save(p) + reg2 = Template_Registration2.load(p, ddevice="cpu") + out = reg2.transform_nii(atlas_vert) + self.assertEqual(out.shape, vert.shape) + + @unittest.skipUnless(_HAS_ELASTIC, "elasticdeform not installed") + def test_deformable_stage_improves_over_rigid_only(self): + """With a truly deformed atlas the deformable stage must reduce the label + mismatch relative to the ``only_rigid=True`` baseline. + + Fixed sigma/points/seed keep this deterministic; ``max_steps`` is picked + so the assertion holds with margin (verified empirically) without the + test becoming slow. + """ + # Use a slightly larger max_steps so convergence is stable, but still cheap. + reg, vert, atlas_vert, _ = self._make(with_pre=True, crop=False, deform=True, max_steps=25) + target_mask = (vert.get_array() != 0).astype(np.int8) + rigid_only = reg.transform_nii(atlas_vert, only_rigid=True) + full = reg.transform_nii(atlas_vert) + rigid_mask = (rigid_only.get_array() != 0).astype(np.int8) + full_mask = (full.get_array() != 0).astype(np.int8) + + def _iou(a, b): + inter = int((a & b).sum()) + union = int((a | b).sum()) + return inter / max(union, 1) + + iou_rigid = _iou(target_mask, rigid_mask) + iou_full = _iou(target_mask, full_mask) + # Deformable stage must not regress overlap. Small positive delta is fine; + # keep the margin loose so tiny numerical wobbles don't fail the test. + self.assertGreaterEqual(iou_full, iou_rigid - 1e-3, f"iou rigid={iou_rigid:.3f} full={iou_full:.3f}") + # And the atlas really was deformed - rigid-only IoU is well below 1. + self.assertLess(iou_rigid, 0.95, f"atlas seems undeformed; iou_rigid={iou_rigid:.3f}") + + +@unittest.skipUnless(_HAS_DEEPALI, "hf-deepali not installed") +class TestFlipHelper(unittest.TestCase): + """The R-axis flip is shared between Template_Registration(2) __init__/warp paths.""" + + def test_nii_flip_is_involution(self): + from TPTBox.registration._deformable.multilabel_segmentation import _flip_r_axis # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + ct, _, _, _ = get_test_ct() + once = _flip_r_axis(ct.copy()) + twice = _flip_r_axis(once) + np.testing.assert_array_equal(twice.get_array(), ct.get_array()) + # And a single flip is NOT the identity on a non-symmetric image. + self.assertFalse(np.array_equal(once.get_array(), ct.get_array())) + + def test_poi_flip_is_involution_and_mirrors_r_axis(self): + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.registration._deformable.multilabel_segmentation import _flip_r_axis # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + _, subreg, vert, _ = get_test_ct() + poi = calc_poi_from_subreg_vert(vert, subreg, subreg_id=[Location.Vertebra_Corpus]).extract_subregion(Location.Vertebra_Corpus) + axis = poi.get_axis("R") + flipped = _flip_r_axis(poi.copy()) + back = _flip_r_axis(flipped) + for k in poi.keys(): + np.testing.assert_allclose(np.array(back[k]), np.array(poi[k]), atol=1e-6) + # Only the R-axis coordinate changed. + orig = np.array(poi[k]) + new = np.array(flipped[k]) + expected = orig.copy() + expected[axis] = poi.shape[axis] - 1 - expected[axis] + np.testing.assert_allclose(new, expected, atol=1e-6) + + def test_poi_global_rejected(self): + """POI_Global has no shape/axis; the helper must refuse instead of returning garbage.""" + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.core.poi_fun.poi_global import POI_Global # noqa: PLC0415 + from TPTBox.registration._deformable.multilabel_segmentation import _flip_r_axis # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + _, subreg, vert, _ = get_test_ct() + poi = calc_poi_from_subreg_vert(vert, subreg, subreg_id=[Location.Vertebra_Corpus]).extract_subregion(Location.Vertebra_Corpus) + with self.assertRaises(TypeError): + _flip_r_axis(POI_Global(poi)) + + +@unittest.skipUnless(_HAS_DEEPALI, "hf-deepali not installed") +class TestSpeedAndMemory(unittest.TestCase): + """Soft speed/memory checks – kept quick enough for CI, no strict deadlines.""" + + def test_deepali_vs_sitk_point_registration(self): + from TPTBox import Location, calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.registration import ( # noqa: PLC0415 + Deepali_Point_Registration, + Point_Registration, + ) + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + ct, subreg, vert, _ = get_test_ct() + poi = calc_poi_from_subreg_vert( + vert, + subreg, + subreg_id=[Location.Vertebra_Corpus, Location.Spinosus_Process, Location.Arcus_Vertebrae], + ).extract_subregion( + Location.Vertebra_Corpus, + Location.Spinosus_Process, + Location.Arcus_Vertebrae, + ) + moving = poi.copy() + moving.origin = tuple(np.asarray(moving.origin) + np.array([2.5, -1.0, 3.0])) + img_moving = ct.copy() + img_moving.origin = moving.origin + + # --- fit time ------------------------------------------------------- + t0 = time.perf_counter() + sitk_reg = Point_Registration(poi, moving, verbose=False) + t_sitk_fit = time.perf_counter() - t0 + t0 = time.perf_counter() + deep_reg = Deepali_Point_Registration(poi, moving, verbose=False, ddevice="cpu") + t_deep_fit = time.perf_counter() - t0 + + # --- warp time ------------------------------------------------------ + t0 = time.perf_counter() + sitk_out = sitk_reg.transform_nii(img_moving) + t_sitk_warp = time.perf_counter() - t0 + tracemalloc.start() + t0 = time.perf_counter() + deep_out = deep_reg.transform_nii(img_moving) + t_deep_warp = time.perf_counter() - t0 + _, peak_deep = tracemalloc.get_traced_memory() + tracemalloc.stop() + + # accuracy comparison against ground truth + orig = ct.get_array().astype(np.float32) + err_sitk = float(np.mean(np.abs(orig - sitk_out.get_array().astype(np.float32)))) + err_deep = float(np.mean(np.abs(orig - deep_out.get_array().astype(np.float32)))) + + # Log summary line so the CI output records the numbers. + print( + f"\n[bench] fit: sitk={t_sitk_fit * 1000:.1f}ms deep={t_deep_fit * 1000:.1f}ms | " + f"warp: sitk={t_sitk_warp * 1000:.1f}ms deep={t_deep_warp * 1000:.1f}ms | " + f"mean-err: sitk={err_sitk:.3f} deep={err_deep:.3f} | " + f"deepali peak mem={peak_deep / (1024 * 1024):.1f} MiB" + ) + # Sanity bounds - accuracy: deepali is typically much better than SITK's + # BSplineResampler here, but we only assert deepali is not massively worse. + self.assertLess(err_deep, max(1.0, err_sitk * 2 + 1.0)) + # Memory: peak within an order of magnitude of the raw volume (~64 MiB for a + # 73^3 float32) - guards against runaway allocations. + raw_mib = orig.nbytes / (1024 * 1024) + self.assertLess(peak_deep / (1024 * 1024), max(raw_mib * 20, 256)) + + +class TestOptionalDeepaliStubs(unittest.TestCase): + """When ``hf-deepali`` isn't installed the deepali-backed entry points + must still *import* – they should raise a helpful ``ImportError`` only + when actually used, and the error message must mention that PyTorch is a + prerequisite too. + """ + + def test_stub_factory_message(self): + import TPTBox.registration as _reg_init # noqa: PLC0415 + + stub_cls = _reg_init._make_missing_deepali_stub("Foo", ImportError("no module named 'deepali'")) + with self.assertRaises(ImportError) as ctx: + stub_cls() + msg = str(ctx.exception) + self.assertIn("Foo", msg) + self.assertIn("hf-deepali", msg) + self.assertIn("PyTorch", msg) + self.assertIn("pip install torch hf-deepali", msg) + + stub_fn = _reg_init._make_missing_deepali_func("bar", ImportError("boom")) + with self.assertRaises(ImportError) as ctx: + stub_fn(1, 2) + self.assertIn("bar()", str(ctx.exception)) + self.assertIn("pip install torch hf-deepali", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main(verbosity=2)