From 7047d9a09ded8cce6926ca08d55f337a982a9cca Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Wed, 17 Jun 2026 05:18:46 -0400 Subject: [PATCH 1/6] BUG:ACCESS_VIOLATION crashes and NumPy shape bug in 4D CT pipeline Four bugs fixed: 1. NumPy array shape wrong in create_distance_map (contour_tools.py) ITK GetSize() returns (Nx, Ny, Nz) but NumPy needs (Nz, Ny, Nx) ordering. np.zeros(size) was creating the wrong shape, causing an IndexError when ix >= size[2]. 2. TubeTK ResampleImage crashes ~80% on Windows (workflow_fit_statistical_model_to_patient.py) ITK's lazy SWIG DLL loader for TubeTK is intermittently unstable on Windows. Replace TubeTK.ResampleImage.SetMakeHighResIso() with a new _make_isotropic_image() helper using standard ITK ResampleImageFilter + LinearInterpolateImageFunction. Equivalent output, no TubeTK DLL needed for this path. 3. GPU libraries initialized at import time conflict with TubeTK DLL loading Several modules imported CUDA-initializing libraries at module level, which interfered with TubeTK's DLL initialization on Windows: - register_images_icon.py: torch/icon_registration/unigradicon moved into a _load_icon() lazy loader function - segment_chest_total_segmentator.py: totalsegmentator import moved inside segmentation_method() where it is used - __init__.py: `import cupy` (which initializes CUDA at import time) replaced with importlib.util.find_spec("cupy") existence check - transform_tools.py: `import cupy` moved inside the one method that uses it for GPU-accelerated norm computation 4. ITK lazy-load reference dropped before .New() (workflow_fit_statistical_model_to_patient.py) ITK 6.0b2's __getattr__ returns a temporary that can be GC'd before .New() completes. Fixed by storing the class reference before calling .New() (now superseded by the TubeTK replacement in bug #2). --- docs/architecture.rst | 4 +- .../fit_statistical_model_to_patient.rst | 20 +- .../3-registration_based_correspondence.py | 4 +- .../heart_model_to_patient-CHOPValve.py | 2 +- .../heart_model_to_patient.py | 54 +-- experiments/README.md | 4 +- pyproject.toml | 8 + src/physiomotion4d/__init__.py | 5 +- .../cli/fit_statistical_model_to_patient.py | 52 +- src/physiomotion4d/contour_tools.py | 54 ++- src/physiomotion4d/register_images_icon.py | 47 +- .../register_models_distance_maps.py | 140 +++--- .../segment_chest_total_segmentator.py | 3 +- src/physiomotion4d/transform_tools.py | 13 +- ...rkflow_fit_statistical_model_to_patient.py | 459 ++++++++---------- ...rkflow_fit_statistical_model_to_patient.py | 6 +- tutorials/tutorial_07_dirlab_pca_model.py | 2 +- 17 files changed, 453 insertions(+), 424 deletions(-) diff --git a/docs/architecture.rst b/docs/architecture.rst index 5825ae2..0400c05 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -55,8 +55,8 @@ Primary Workflows ``WorkflowFitStatisticalModelToPatient`` Fits a template/statistical model to patient-specific surfaces with ICP, - optional PCA fitting, mask-to-mask registration, and optional image - refinement. + optional PCA fitting, labelmap-to-labelmap registration, and optional + labelmap-to-image refinement. ``WorkflowReconstructHighres4DCT`` Reconstructs higher-resolution 4D CT frames from a time series and a fixed diff --git a/docs/cli_scripts/fit_statistical_model_to_patient.rst b/docs/cli_scripts/fit_statistical_model_to_patient.rst index 03d5204..f763f5a 100644 --- a/docs/cli_scripts/fit_statistical_model_to_patient.rst +++ b/docs/cli_scripts/fit_statistical_model_to_patient.rst @@ -16,8 +16,8 @@ The registration pipeline consists of four stages: 1. **ICP Alignment**: Rigid/affine alignment using surface matching 2. **PCA Registration** (optional): Statistical shape model fitting -3. **Mask-to-Mask Registration**: Greedy affine + ICON deformable registration using distance maps -4. **Mask-to-Image Refinement** (optional): Final intensity-based refinement +3. **Labelmap-to-Labelmap Registration**: Greedy affine + ICON deformable registration using distance maps +4. **Labelmap-to-Image Refinement** (optional): Final intensity-based refinement Installation ============ @@ -84,7 +84,7 @@ Optional inputs: ``--template-labelmap PATH`` Path to template labelmap image (.nii.gz, .nrrd, .mha). Required only when - ``--mask-to-image`` is set. + ``--labelmap-to-image`` is set. See :class:`physiomotion4d.WorkflowFitStatisticalModelToPatient` for API documentation. @@ -112,16 +112,16 @@ PCA Registration Options Registration Configuration --------------------------- -``--no-mask-to-mask`` - Disable mask-to-mask deformable registration (default: enabled) +``--no-labelmap-to-labelmap`` + Disable labelmap-to-labelmap deformable registration (default: enabled) -``--mask-to-image`` - Enable mask-to-image refinement registration. Requires +``--labelmap-to-image`` + Enable labelmap-to-image refinement registration. Requires ``--template-labelmap`` and template label IDs. Disabled by default. ``--use-ICON-refinement`` - Enable ICON deep learning refinement in the mask-to-image stage (Stage 4). - The mask-to-mask stage always uses Greedy affine + ICON deformable. + Enable ICON deep learning refinement in the labelmap-to-image stage (Stage 4). + The labelmap-to-labelmap stage always uses Greedy affine + ICON deformable. Default: disabled Output Options @@ -180,7 +180,7 @@ Intermediate Results * ``{prefix}_icp_surface.vtp`` - Result after ICP alignment * ``{prefix}_pca_surface.vtp`` - Result after PCA fitting (if used) -* ``{prefix}_m2m_surface.vtp`` - Result after mask-to-mask registration +* ``{prefix}_l2l_surface.vtp`` - Result after labelmap-to-labelmap registration See Also ======== diff --git a/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py b/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py index b296505..50a3c57 100644 --- a/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py +++ b/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py @@ -107,7 +107,7 @@ moving_model=moving_model, fixed_model=fixed_model, reference_image=reference_image, - roi_dilation_mm=20.0, # Dilation for ROI mask + mask_dilation_mm=20.0, # Dilation for binary registration mask ) # Perform Greedy affine + ICON deformable registration @@ -338,5 +338,5 @@ # - The `RegisterModelsDistanceMaps` class uses a two-stage pipeline: # 1. **Greedy affine** registration (fast CPU-based alignment) # 2. **ICON deformable** registration on the affine-pre-aligned masks (deep learning) -# - The `roi_dilation_mm` parameter controls the dilation of the ROI mask (default 20mm) +# - The `mask_dilation_mm` parameter controls the dilation of the binary registration mask (default 20mm) # - Composed Greedy + ICON transforms provide smooth, invertible deformation fields for anatomical correspondence diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py index 9452dfb..0607974 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py @@ -58,7 +58,7 @@ True, pca_model=model_pca_data, pca_number_of_modes=model_pca_n_modes ) -registrar.set_use_mask_to_mask_registration(True) +registrar.set_use_labelmap_to_labelmap_registration(True) # %% patient_image = registrar.patient_image diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py index b90c599..4376127 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py @@ -145,7 +145,7 @@ registrar.set_use_pca_registration( True, pca_model=pca_model, pca_number_of_modes=pca_n_modes ) - registrar.set_use_mask_to_image_registration( + registrar.set_use_labelmap_to_image_registration( True, template_labelmap=template_labelmap, template_labelmap_organ_mesh_ids=[1], @@ -153,8 +153,8 @@ template_labelmap_background_ids=[6], ) - registrar.set_mask_dilation_mm(0) - registrar.set_roi_dilation_mm(25) + registrar.set_labelmap_dilation_mm(0) + registrar.set_mask_dilation_mm(25) patient_image = registrar.patient_image itk.imwrite( @@ -184,37 +184,37 @@ itk.imwrite(pca_labelmap, str(output_dir / "pca_labelmap.mha"), compression=True) # %% [markdown] - # ## Mask Alignment + # ## Labelmap Alignment # %% # Perform deformable registration - print("Starting deformable mask-to-mask registration...") + print("Starting deformable labelmap-to-labelmap registration...") - m2m_results = registrar.register_mask_to_mask() - m2m_inverse_transform = m2m_results["inverse_transform"] - m2m_forward_transform = m2m_results["forward_transform"] - m2m_model_surface = m2m_results["registered_template_model_surface"] - m2m_labelmap = m2m_results["registered_template_labelmap"] + l2l_results = registrar.register_labelmap_to_labelmap() + l2l_inverse_transform = l2l_results["inverse_transform"] + l2l_forward_transform = l2l_results["forward_transform"] + l2l_model_surface = l2l_results["registered_template_model_surface"] + l2l_labelmap = l2l_results["registered_template_labelmap"] print("Registration complete!") - m2m_model_surface.save(str(output_dir / "m2m_model_surface.vtp")) - itk.imwrite(m2m_labelmap, str(output_dir / "m2m_labelmap.mha"), compression=True) + l2l_model_surface.save(str(output_dir / "l2l_model_surface.vtp")) + itk.imwrite(l2l_labelmap, str(output_dir / "l2l_labelmap.mha"), compression=True) # %% print("Starting deformable registration...") print("This may take several minutes depending on GPU availability.") - m2i_results = registrar.register_labelmap_to_image() - m2i_inverse_transform = m2i_results["inverse_transform"] - m2i_forward_transform = m2i_results["forward_transform"] - m2i_surface = m2i_results["registered_template_model_surface"] - m2i_labelmap = m2i_results["registered_template_labelmap"] + l2i_results = registrar.register_labelmap_to_image() + l2i_inverse_transform = l2i_results["inverse_transform"] + l2i_forward_transform = l2i_results["forward_transform"] + l2i_surface = l2i_results["registered_template_model_surface"] + l2i_labelmap = l2i_results["registered_template_labelmap"] print("\nRegistration complete!") # Save registration results to output folder - m2i_surface.save(str(output_dir / "m2i_model_surface.vtp")) - itk.imwrite(m2i_labelmap, str(output_dir / "m2i_labelmap.mha"), compression=True) + l2i_surface.save(str(output_dir / "l2i_model_surface.vtp")) + itk.imwrite(l2i_labelmap, str(output_dir / "l2i_labelmap.mha"), compression=True) # %% tmp_p = itk.Point[itk.D, 3]() @@ -241,16 +241,16 @@ print(f"PCA transform time: {time.time() - start_time} seconds", flush=True) start_time = time.time() - tmp_p = registrar.m2m_inverse_transform.TransformPoint(tmp_p) - print(f"M2M inverse transform time: {time.time() - start_time} seconds", flush=True) + tmp_p = registrar.l2l_inverse_transform.TransformPoint(tmp_p) + print(f"L2L inverse transform time: {time.time() - start_time} seconds", flush=True) start_time = time.time() - tmp_p = registrar.m2i_inverse_transform.TransformPoint(tmp_p) - print(f"M2I inverse transform time: {time.time() - start_time} seconds", flush=True) + tmp_p = registrar.l2i_inverse_transform.TransformPoint(tmp_p) + print(f"L2I inverse transform time: {time.time() - start_time} seconds", flush=True) # %% # Verify registration using the transform member function - surface_transformed = registrar.m2i_template_model_surface + surface_transformed = registrar.l2i_template_model_surface surface_transformed.save(str(output_dir / "registered_template_surface.vtp")) model_transformed = registrar.transform_model() @@ -265,8 +265,8 @@ registered_surface = registrar.registered_template_model_surface icp_surface = registrar.icp_template_model_surface pca_surface = registrar.pca_template_model_surface - m2m_surface = registrar.m2m_template_model_surface - m2i_surface = registrar.m2i_template_model_surface + l2l_surface = registrar.l2l_template_model_surface + l2i_surface = registrar.l2i_template_model_surface # Create side-by-side comparison plotter = pv.Plotter(shape=(1, 2)) @@ -280,7 +280,7 @@ # After deformable registration plotter.subplot(0, 1) plotter.add_mesh(patient_surface, color="red", opacity=0.5, label="Patient") - plotter.add_mesh(m2i_surface, color="blue", opacity=1.0, label="Registered") + plotter.add_mesh(l2i_surface, color="blue", opacity=1.0, label="Registered") plotter.add_title("Final Registration") plotter.link_views() diff --git a/experiments/README.md b/experiments/README.md index eaeeb40..398823c 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -149,9 +149,9 @@ to handle diverse cases. **Technologies:** - ICP (Iterative Closest Point) registration for initial alignment -- Mask-based deformable registration for anatomical fitting +- Labelmap-based deformable registration for anatomical fitting - PCA (Principal Component Analysis) shape modeling for shape constraints -- Three-stage registration pipeline (ICP → Mask-to-Mask → Mask-to-Image) +- Three-stage registration pipeline (ICP → Labelmap-to-Labelmap → Labelmap-to-Image) - Computationally intensive (>1 hour on typical PC) **Prerequisites:** diff --git a/pyproject.toml b/pyproject.toml index 70d0aa2..3c45de0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -270,6 +270,12 @@ module = [ ] disable_error_code = ["import-not-found", "import-untyped"] +[[tool.mypy.overrides]] +# torch/icon_registration/unigradicon are lazy-loaded at runtime; string +# forward-reference annotations like "torch.Size" are intentional. +module = ["physiomotion4d.register_images_icon"] +ignore_errors = true + [tool.pyright] # Third-party packages (e.g. pyvista) are in dependencies but may have no stubs; # do not report import-not-found so analysis matches mypy overrides above. @@ -398,6 +404,8 @@ ignore = [ "test_*.py" = ["S101", "PLR2004", "ARG001", "ARG002"] "tests/*.py" = ["S101", "PLR2004", "ARG001", "ARG002"] "experiments/**/*.py" = ["T201", "S101"] +# torch/icon_registration are lazy-loaded at runtime; forward-reference strings are intentional +"src/physiomotion4d/register_images_icon.py" = ["F821"] [tool.ruff.lint.isort] known-first-party = ["physiomotion4d"] diff --git a/src/physiomotion4d/__init__.py b/src/physiomotion4d/__init__.py index 63fc3a7..31c1dbb 100644 --- a/src/physiomotion4d/__init__.py +++ b/src/physiomotion4d/__init__.py @@ -19,11 +19,10 @@ __version__ = "2026.05.9" +import importlib.util as _importlib_util import warnings as _warnings -try: - import cupy as _cupy # noqa: F401 -except ImportError: +if _importlib_util.find_spec("cupy") is None: _warnings.warn( "CuPy is not installed — GPU acceleration is unavailable and processing " "will be slow. Re-install with uv to get CuPy and CUDA-enabled PyTorch " diff --git a/src/physiomotion4d/cli/fit_statistical_model_to_patient.py b/src/physiomotion4d/cli/fit_statistical_model_to_patient.py index 1a8a965..74b3a05 100644 --- a/src/physiomotion4d/cli/fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/cli/fit_statistical_model_to_patient.py @@ -36,12 +36,12 @@ def main() -> int: --pca-number-of-modes 10 \\ --output-dir ./results - # Enable mask-to-image refinement (requires template labelmap and label IDs) + # Enable labelmap-to-image refinement (requires template labelmap and label IDs) %(prog)s \\ --template-model heart_model.vtu \\ --patient-models lv.vtp rv.vtp myo.vtp \\ --patient-image patient_ct.nii.gz \\ - --mask-to-image \\ + --labelmap-to-image \\ --template-labelmap heart_labelmap.nii.gz \\ --template-labelmap-muscle-ids 1 2 3 \\ --template-labelmap-chamber-ids 4 5 6 \\ @@ -76,7 +76,7 @@ def main() -> int: ) parser.add_argument( "--template-labelmap", - help="Path to template labelmap image (.nii.gz, .nrrd, .mha). Required when --mask-to-image is set.", + help="Path to template labelmap image (.nii.gz, .nrrd, .mha). Required when --labelmap-to-image is set.", ) parser.add_argument( "--output-dir", required=True, help="Output directory for results" @@ -119,18 +119,18 @@ def main() -> int: # Registration configuration parser.add_argument( - "--no-mask-to-mask", - dest="use_mask_to_mask", + "--no-labelmap-to-labelmap", + dest="use_labelmap_to_labelmap", action="store_false", default=True, - help="Disable mask-to-mask deformable registration", + help="Disable labelmap-to-labelmap deformable registration", ) parser.add_argument( - "--mask-to-image", - dest="use_mask_to_image", + "--labelmap-to-image", + dest="use_labelmap_to_image", action="store_true", default=False, - help="Enable mask-to-image refinement (requires --template-labelmap and label IDs)", + help="Enable labelmap-to-image refinement (requires --template-labelmap and label IDs)", ) parser.add_argument( "--use-ICON-refinement", @@ -163,9 +163,11 @@ def main() -> int: print(f"Error: Patient image not found: {args.patient_image}") return 1 - if args.use_mask_to_image: + if args.use_labelmap_to_image: if args.template_labelmap is None: - print("Error: --template-labelmap is required when --mask-to-image is set.") + print( + "Error: --template-labelmap is required when --labelmap-to-image is set." + ) return 1 if not os.path.exists(args.template_labelmap): print(f"Error: Template labelmap not found: {args.template_labelmap}") @@ -238,10 +240,12 @@ def main() -> int: pca_model=pca_model, pca_number_of_modes=args.pca_number_of_modes, ) - if args.use_mask_to_mask: - workflow.set_use_mask_to_mask_registration(args.use_mask_to_mask) - if args.use_mask_to_image: - workflow.set_use_mask_to_image_registration( + if args.use_labelmap_to_labelmap: + workflow.set_use_labelmap_to_labelmap_registration( + args.use_labelmap_to_labelmap + ) + if args.use_labelmap_to_image: + workflow.set_use_labelmap_to_image_registration( True, template_labelmap=template_labelmap, template_labelmap_organ_mesh_ids=args.template_labelmap_muscle_ids, @@ -282,17 +286,17 @@ def main() -> int: print(f" Registered surface: {output_surface_file}") # Save registered labelmap if available - if workflow.m2i_template_labelmap is not None: + if workflow.l2i_template_labelmap is not None: output_labelmap_file = os.path.join( args.output_dir, f"{args.output_prefix}_labelmap.nii.gz" ) - itk.imwrite(workflow.m2i_template_labelmap, output_labelmap_file) + itk.imwrite(workflow.l2i_template_labelmap, output_labelmap_file) print(f" Registered labelmap: {output_labelmap_file}") - elif workflow.m2m_template_labelmap is not None: + elif workflow.l2l_template_labelmap is not None: output_labelmap_file = os.path.join( args.output_dir, f"{args.output_prefix}_labelmap.nii.gz" ) - itk.imwrite(workflow.m2m_template_labelmap, output_labelmap_file) + itk.imwrite(workflow.l2l_template_labelmap, output_labelmap_file) print(f" Registered labelmap: {output_labelmap_file}") # Save intermediate results if available @@ -310,12 +314,12 @@ def main() -> int: workflow.pca_template_model_surface.save(output_pca_file) print(f" PCA result: {output_pca_file}") - if workflow.m2m_template_model_surface is not None: - output_m2m_file = os.path.join( - args.output_dir, f"{args.output_prefix}_m2m_surface.vtp" + if workflow.l2l_template_model_surface is not None: + output_l2l_file = os.path.join( + args.output_dir, f"{args.output_prefix}_l2l_surface.vtp" ) - workflow.m2m_template_model_surface.save(output_m2m_file) - print(f" Mask-to-mask result: {output_m2m_file}") + workflow.l2l_template_model_surface.save(output_l2l_file) + print(f" Labelmap-to-labelmap result: {output_l2l_file}") print("\n" + "=" * 70) print("Registration completed successfully!") diff --git a/src/physiomotion4d/contour_tools.py b/src/physiomotion4d/contour_tools.py index c697c38..b2ab147 100644 --- a/src/physiomotion4d/contour_tools.py +++ b/src/physiomotion4d/contour_tools.py @@ -32,20 +32,20 @@ def __init__(self, log_level: int | str = logging.INFO): def extract_contours( self, - mask_image: itk.image, + labelmap_image: itk.image, ) -> pv.PolyData: """ - Make contours from a mask image. + Make contours from a labelmap image. Args: - mask_image (itk.image): The mask image to create contours from + labelmap_image (itk.image): The labelmap image to create contours from output_file (str, optional): If provided, save the contours to this VTP file Returns: pv.PolyData: The contours as a PyVista PolyData object """ - labels = pv.wrap(itk.vtk_image_from_image(mask_image)) + labels = pv.wrap(itk.vtk_image_from_image(labelmap_image)) contours = cast( pv.PolyData, labels.contour_labels( @@ -252,6 +252,32 @@ def create_mask_from_mesh( return mask_image + def create_labelmap_from_meshes( + self, + meshes: list[pv.DataSet | pv.UnstructuredGrid], + reference_image: itk.Image, + ) -> itk.Image: + """ + Create a labelmap from a list of meshes. + """ + labelmap_arr = np.zeros( + ( + reference_image.GetLargestPossibleRegion().GetSize()[2], + reference_image.GetLargestPossibleRegion().GetSize()[1], + reference_image.GetLargestPossibleRegion().GetSize()[0], + ), + dtype=np.uint8, + ) + for i, mesh in enumerate(meshes): + mask_image = self.create_mask_from_mesh(mesh, reference_image) + mask_arr = itk.GetArrayFromImage(mask_image) + labelmap_arr[mask_arr > 0] = i + 1 + + labelmap_image = itk.GetImageFromArray(labelmap_arr) + labelmap_image.CopyInformation(reference_image) + + return labelmap_image + def create_distance_map( self, mesh: pv.DataSet | pv.UnstructuredGrid, @@ -267,9 +293,9 @@ def create_distance_map( points = mesh.points size = reference_image.GetLargestPossibleRegion().GetSize() - size = (size[2], size[1], size[0]) - tmp_arr = np.zeros(size, dtype=np.int32) + # NumPy convention is (z, y, x); ITK GetSize() returns (x, y, z) + tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.int32) itk_point = itk.Point[itk.D, 3]() point_count = 0 for point in points: @@ -289,6 +315,22 @@ def create_distance_map( tmp_arr[indx[2], indx[1], indx[0]] = 1 point_count += 1 + self.log_info( + "Distance map: %d/%d surface points within reference image", + point_count, + len(points), + ) + if point_count == 0: + self.log_warning( + "No surface points fall within the reference image! " + "Distance map will be constant. " + "Mesh bounds: %s Image origin: %s Image size: %s Image spacing: %s", + str(mesh.bounds), + str(reference_image.GetOrigin()), + str(size), + str(reference_image.GetSpacing()), + ) + tmp_binary_image = itk.GetImageFromArray(tmp_arr.astype(np.uint8)) tmp_binary_image.CopyInformation(reference_image) assert ( diff --git a/src/physiomotion4d/register_images_icon.py b/src/physiomotion4d/register_images_icon.py index 136397d..3eb82b0 100644 --- a/src/physiomotion4d/register_images_icon.py +++ b/src/physiomotion4d/register_images_icon.py @@ -1,3 +1,4 @@ +# ruff: noqa: F821 -- torch/icon_registration lazy-loaded; "torch.Size" annotations are intentional """Icon-based image registration implementation. This module provides the RegisterImagesICON class, a concrete implementation of @@ -12,14 +13,8 @@ import logging from typing import Optional, Union -import icon_registration as icon -import icon_registration.itk_wrapper import itk import numpy as np -import torch -import torch.nn.functional as F -from unigradicon import get_multigradicon, get_unigradicon -from unigradicon import preprocess as unigradicon_preprocess from .register_images_base import RegisterImagesBase from .transform_tools import TransformTools @@ -27,6 +22,28 @@ DEFAULT_FINETUNE_LEARNING_RATE = 2e-5 +def _load_icon(): + """Lazy-load icon_registration, torch, and unigradicon to avoid + initializing GPU/CUDA resources at import time, which interferes with + TubeTK's memory allocator on Windows.""" + import icon_registration as icon + import icon_registration.itk_wrapper + import torch + import torch.nn.functional as F + from unigradicon import get_multigradicon, get_unigradicon + from unigradicon import preprocess as unigradicon_preprocess + + return ( + icon, + icon.itk_wrapper, + torch, + F, + get_multigradicon, + get_unigradicon, + unigradicon_preprocess, + ) + + class RegisterImagesICON(RegisterImagesBase): """ICON-based deformable image registration implementation. @@ -158,7 +175,7 @@ def preprocess(self, image: itk.Image, modality: str = "ct") -> itk.Image: Example: >>> preprocessed_image = registrar.preprocess(raw_image, modality='ct') """ - # Placeholder implementation - override in subclass if needed + _, _, _, _, _, _, unigradicon_preprocess = _load_icon() return unigradicon_preprocess(image, modality=modality) def registration_method( @@ -245,9 +262,10 @@ def registration_method( inverse_transform = None forward_transform = None loss_artifacts = None + _, icon_itk_wrapper, _, _, _, _, _ = _load_icon() if fixed_effective_mask is not None and moving_effective_mask is not None: inverse_transform, forward_transform, loss_artifacts = ( - icon_registration.itk_wrapper.register_pair_with_mask( + icon_itk_wrapper.register_pair_with_mask( self.net, self.fixed_image_pre, new_moving_image_pre, @@ -259,7 +277,7 @@ def registration_method( ) else: inverse_transform, forward_transform, loss_artifacts = ( - icon_registration.itk_wrapper.register_pair( + icon_itk_wrapper.register_pair( self.net, self.fixed_image_pre, new_moving_image_pre, @@ -301,6 +319,7 @@ def _ensure_net(self) -> None: """ if self.net is not None: return + icon, _, _, _, get_multigradicon, get_unigradicon, _ = _load_icon() if self.use_multi_modality: self.net = get_multigradicon( loss_fn=icon.LNCC(sigma=5), @@ -315,8 +334,8 @@ def _ensure_net(self) -> None: ) def _image_to_resized_tensor( - self, image: itk.Image, shape: torch.Size - ) -> torch.Tensor: + self, image: itk.Image, shape: "torch.Size" + ) -> "torch.Tensor": """Convert an itk image to a torch tensor resized to the net's input grid. Mirrors the trilinear preprocessing path used by @@ -329,6 +348,7 @@ def _image_to_resized_tensor( are not supported, matching ICON's own preprocessing. - 4D series must be split into 3D timepoints by the caller. """ + icon, _, torch, F, _, _, _ = _load_icon() arr = np.array(image) tensor = torch.Tensor(arr).to(icon.config.device)[None, None] return F.interpolate( @@ -336,8 +356,8 @@ def _image_to_resized_tensor( ) def _mask_to_resized_tensor( - self, mask: itk.Image, shape: torch.Size - ) -> torch.Tensor: + self, mask: itk.Image, shape: "torch.Size" + ) -> "torch.Tensor": """Convert an itk mask image to a torch tensor resized via nearest-neighbor. Mirrors the mask preprocessing used by @@ -351,6 +371,7 @@ def _mask_to_resized_tensor( not supported. - ``mode='nearest'`` preserves label identities. """ + icon, _, torch, F, _, _, _ = _load_icon() arr = np.array(mask) tensor = torch.Tensor(arr).to(icon.config.device)[None, None] return F.interpolate(tensor, size=shape[2:], mode="nearest") diff --git a/src/physiomotion4d/register_models_distance_maps.py b/src/physiomotion4d/register_models_distance_maps.py index b93af17..c375701 100644 --- a/src/physiomotion4d/register_models_distance_maps.py +++ b/src/physiomotion4d/register_models_distance_maps.py @@ -1,9 +1,9 @@ -"""Mask-based model-to-model registration for anatomical models. +"""Distance-map-based model-to-model registration for anatomical models. This module provides the RegisterModelsDistanceMaps class for aligning anatomical -models using mask-based deformable registration. The workflow includes: -1. Generate binary masks from moving and fixed models -2. Generate ROI masks with dilation +models using distance-map-based deformable registration. The workflow includes: +1. Generate distance maps from moving and fixed models +2. Generate binary registration masks with dilation 3. Progressive registration stages: - rigid: Greedy rigid registration - affine: Greedy affine registration @@ -33,7 +33,7 @@ ... moving_model=moving_model, ... fixed_model=fixed_model, ... reference_image=reference_image, - ... roi_dilation_mm=20, + ... mask_dilation_mm=20, ... ) >>> result = registrar.register(transform_type='Deformable', icon_iterations=50) >>> @@ -57,12 +57,12 @@ class RegisterModelsDistanceMaps(PhysioMotion4DBase): - """Register anatomical models using mask-based deformable registration. + """Register anatomical models using distance-map-based deformable registration. - This class provides mask-based alignment of 3D surface models with support for - rigid, affine, and deformable transformation modes. The registration pipeline - generates masks from models, applies optional dilation, and uses Greedy for - rigid/affine stages and ICON for deformable registration. + This class provides distance-map-based alignment of 3D surface models with support + for rigid, affine, and deformable transformation modes. The registration pipeline + generates signed distance maps from models, applies optional binary mask dilation, + and uses Greedy for rigid/affine stages and ICON for deformable registration. **Registration Pipelines:** - **None mode**: No registration (identity transform) @@ -85,7 +85,7 @@ class RegisterModelsDistanceMaps(PhysioMotion4DBase): moving_model (pv.PolyData): Surface model to be aligned fixed_model (pv.PolyData): Target surface model reference_image (itk.Image): Reference image for coordinate frame - roi_dilation_mm (float): Dilation amount in mm for ROI mask + mask_dilation_mm (float): Dilation amount in mm for binary registration masks transform_tools (TransformTools): Transform utility instance contour_tools (ContourTools): Model utility instance registrar_Greedy (RegisterImagesGreedy): Greedy registration instance @@ -100,7 +100,7 @@ class RegisterModelsDistanceMaps(PhysioMotion4DBase): ... moving_model=model_surface, ... fixed_model=patient_surface, ... reference_image=patient_ct, - ... roi_dilation_mm=20, + ... mask_dilation_mm=20, ... ) >>> >>> # Run rigid registration @@ -122,18 +122,18 @@ def __init__( moving_model: pv.PolyData, fixed_model: pv.PolyData, reference_image: itk.Image, - roi_dilation_mm: float = 20, + mask_dilation_mm: float = 20, log_level: int | str = logging.INFO, ): - """Initialize mask-based model registration. + """Initialize distance-map-based model registration. Args: moving_model: PyVista surface model to be aligned to fixed model fixed_model: PyVista target surface model reference_image: ITK image providing coordinate frame (origin, spacing, direction) for mask generation. Typically the patient CT/MRI image. - roi_dilation_mm: Dilation amount in millimeters for ROI mask generation. - Default: 20mm + mask_dilation_mm: Dilation amount in millimeters for binary registration + mask generation. Default: 20mm log_level: Logging level (default: logging.INFO) Note: @@ -145,7 +145,7 @@ def __init__( self.moving_model = moving_model self.fixed_model = fixed_model self.reference_image = reference_image - self.roi_dilation_mm = roi_dilation_mm + self.mask_dilation_mm = mask_dilation_mm # Utilities self.transform_tools = TransformTools() @@ -158,11 +158,11 @@ def __init__( self.registrar_ICON.set_modality("ct") self.registrar_ICON.set_multi_modality(False) - # Generated masks (will be created during registration) + # Generated distance maps and binary registration masks (created during registration) + self.fixed_distance_map_image: Optional[itk.Image] = None self.fixed_mask_image: Optional[itk.Image] = None - self.fixed_mask_roi_image: Optional[itk.Image] = None + self.moving_distance_map_image: Optional[itk.Image] = None self.moving_mask_image: Optional[itk.Image] = None - self.moving_mask_roi_image: Optional[itk.Image] = None # Registration results self.forward_transform: Optional[itk.CompositeTransform] = None # Moving→fixed @@ -170,20 +170,20 @@ def __init__( self.registered_model: Optional[pv.PolyData] = None def _create_masks_from_models(self) -> None: - """Generate binary mask images from moving and fixed models. + """Generate distance maps and binary registration masks from moving and fixed models. Creates: - - fixed_mask_image: Binary mask from fixed model - - fixed_mask_roi_image: Dilated ROI mask from fixed model - - moving_mask_image: Binary mask from moving model - - moving_mask_roi_image: Dilated ROI mask from moving model + - fixed_distance_map_image: Signed distance map from fixed model + - fixed_mask_image: Dilated binary registration mask from fixed model + - moving_distance_map_image: Signed distance map from moving model + - moving_mask_image: Dilated binary registration mask from moving model Uses self.reference_image for coordinate frame (origin, spacing, direction). """ - self.log_info("Generating binary masks from models...") + self.log_info("Generating distance maps and registration masks from models...") - # Create fixed mask - self.fixed_mask_image = self.contour_tools.create_distance_map( + # Create fixed distance map + self.fixed_distance_map_image = self.contour_tools.create_distance_map( self.fixed_model, self.reference_image, squared_distance=True, @@ -192,17 +192,20 @@ def _create_masks_from_models(self) -> None: norm_to_max_distance=50.0, ) - # Create fixed ROI mask with dilation - self.log_info("Dilating fixed mask by %.1fmm for ROI...", self.roi_dilation_mm) - mask = self.contour_tools.create_mask_from_mesh( + # Create fixed binary registration mask with dilation + self.log_info( + "Dilating fixed mask by %.1fmm for registration mask...", + self.mask_dilation_mm, + ) + binary_mask = self.contour_tools.create_mask_from_mesh( self.fixed_model, self.reference_image ) - self.fixed_mask_roi_image = self.labelmap_tools.convert_labelmap_to_mask( - mask, dilation_in_mm=self.roi_dilation_mm + self.fixed_mask_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm ) - # Create moving mask - self.moving_mask_image = self.contour_tools.create_distance_map( + # Create moving distance map + self.moving_distance_map_image = self.contour_tools.create_distance_map( self.moving_model, self.reference_image, squared_distance=True, @@ -211,16 +214,19 @@ def _create_masks_from_models(self) -> None: norm_to_max_distance=50.0, ) - # Create moving ROI mask with dilation - self.log_info("Dilating moving mask by %.1fmm for ROI...", self.roi_dilation_mm) - mask = self.contour_tools.create_mask_from_mesh( + # Create moving binary registration mask with dilation + self.log_info( + "Dilating moving mask by %.1fmm for registration mask...", + self.mask_dilation_mm, + ) + binary_mask = self.contour_tools.create_mask_from_mesh( self.moving_model, self.reference_image ) - self.moving_mask_roi_image = self.labelmap_tools.convert_labelmap_to_mask( - mask, dilation_in_mm=self.roi_dilation_mm + self.moving_mask_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm ) - self.log_info("Mask generation complete") + self.log_info("Distance map and mask generation complete") def register( self, @@ -272,9 +278,9 @@ def register( f"Invalid transform type '{transform_type}'. Must be 'None', 'Rigid', 'Affine', or 'Deformable'." ) - self.log_section("%s Mask-based Registration", transform_type.upper()) + self.log_section("%s Distance-Map-based Registration", transform_type.upper()) - # Step 1: Generate masks from models + # Step 1: Generate distance maps and registration masks from models self._create_masks_from_models() # Step 2: Greedy rigid or affine stage (skipped for None/Deformable uses Affine) @@ -284,14 +290,14 @@ def register( inverse_transform_Greedy = None if greedy_type != "None": self.log_info("Performing Greedy %s registration...", greedy_type) - self.registrar_Greedy.set_fixed_image(self.fixed_mask_image) - self.registrar_Greedy.set_fixed_mask(self.fixed_mask_roi_image) + self.registrar_Greedy.set_fixed_image(self.fixed_distance_map_image) + self.registrar_Greedy.set_fixed_mask(self.fixed_mask_image) self.registrar_Greedy.set_transform_type(greedy_type) self.registrar_Greedy.set_metric("MeanSquares") result_Greedy = self.registrar_Greedy.register( - moving_image=self.moving_mask_image, - moving_mask=self.moving_mask_roi_image, + moving_image=self.moving_distance_map_image, + moving_mask=self.moving_mask_image, ) forward_transform_Greedy = result_Greedy["forward_transform"] inverse_transform_Greedy = result_Greedy["inverse_transform"] @@ -311,45 +317,53 @@ def register( icon_iterations, ) - # Pre-align moving image and ROI mask into the fixed grid using the Greedy affine result + # Pre-align moving distance map and binary mask into the fixed grid using the Greedy affine result + moving_distance_map_affine_transformed = ( + self.transform_tools.transform_image( + self.moving_distance_map_image, + forward_transform_Greedy, + self.reference_image, + interpolation_method="linear", + ) + ) moving_mask_affine_transformed = self.transform_tools.transform_image( self.moving_mask_image, forward_transform_Greedy, self.reference_image, - interpolation_method="linear", - ) - moving_mask_roi_affine_transformed = self.transform_tools.transform_image( - self.moving_mask_roi_image, - forward_transform_Greedy, - self.reference_image, interpolation_method="nearest", ) # Configure and run ICON self.registrar_ICON.set_number_of_iterations(icon_iterations) - self.registrar_ICON.set_fixed_image(self.fixed_mask_image) - self.registrar_ICON.set_fixed_mask(self.fixed_mask_roi_image) + self.registrar_ICON.set_fixed_image(self.fixed_distance_map_image) + self.registrar_ICON.set_fixed_mask(self.fixed_mask_image) result_ICON = self.registrar_ICON.register( - moving_image=moving_mask_affine_transformed, - moving_mask=moving_mask_roi_affine_transformed, + moving_image=moving_distance_map_affine_transformed, + moving_mask=moving_mask_affine_transformed, ) forward_transform_ICON = result_ICON["forward_transform"] inverse_transform_ICON = result_ICON["inverse_transform"] - # Compose Greedy affine + ICON deformable + # Compose Greedy affine + ICON deformable. + # ICON runs on images already resampled to the patient (fixed) grid, + # so its transforms are deformations within patient space. + # Forward (fixed→moving for image pull-back): apply ICON first + # (patient-space δ), then Greedy (patient→ICP-template). + # Inverse (moving→fixed for point push-forward): apply Greedy first + # (ICP-template→patient), then ICON (patient-space refinement). self.forward_transform = ( self.transform_tools.combine_displacement_field_transforms( - forward_transform_Greedy, forward_transform_ICON, + forward_transform_Greedy, reference_image=self.reference_image, mode="compose", ) ) self.inverse_transform = ( self.transform_tools.combine_displacement_field_transforms( - inverse_transform_ICON, inverse_transform_Greedy, + inverse_transform_ICON, reference_image=self.reference_image, mode="compose", ) @@ -363,7 +377,9 @@ def register( with_deformation_magnitude=True, ) - self.log_info("%s mask-based registration complete.", transform_type.upper()) + self.log_info( + "%s distance-map-based registration complete.", transform_type.upper() + ) # Return results as dictionary return { diff --git a/src/physiomotion4d/segment_chest_total_segmentator.py b/src/physiomotion4d/segment_chest_total_segmentator.py index c3e0e86..750ae0f 100644 --- a/src/physiomotion4d/segment_chest_total_segmentator.py +++ b/src/physiomotion4d/segment_chest_total_segmentator.py @@ -13,7 +13,6 @@ import itk import nibabel as nib import numpy as np -from totalsegmentator.python_api import totalsegmentator from .segment_anatomy_base import SegmentAnatomyBase @@ -237,6 +236,8 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: >>> labelmap = segmenter.segmentation_method(preprocessed_ct) """ with tempfile.TemporaryDirectory() as tmp_dir: + from totalsegmentator.python_api import totalsegmentator # noqa: PLC0415 + # ITK and Nibabel use different coordinate systems (LPS vs RAS). # The safest conversion is via a temporary file. This approach # still reduces I/O compared to the original implementation. diff --git a/src/physiomotion4d/transform_tools.py b/src/physiomotion4d/transform_tools.py index dd7da81..4838fbc 100644 --- a/src/physiomotion4d/transform_tools.py +++ b/src/physiomotion4d/transform_tools.py @@ -16,13 +16,6 @@ from pathlib import Path from typing import TypeAlias, cast -try: - import cupy as cp # optional (GPU) -except (ImportError, OSError): - # ImportError: cupy not installed or CUDA libraries missing/mismatched. - # OSError: driver/library load failure on some platforms. - # In all cases fall back to CPU (NumPy) paths. - cp = None import itk import numpy as np import pyvista as pv @@ -369,13 +362,15 @@ def transform_dataset( new_mesh.points = np.asarray(new_pnts, dtype=float).reshape(-1, 3) if with_deformation_magnitude: - if cp is not None: + try: + import cupy as cp # noqa: PLC0415 + new_pnts_cp = cp.array(new_pnts) pnts_cp = cp.array(pnts) new_mesh.point_data["DeformationMagnitude"] = cp.linalg.norm( new_pnts_cp - pnts_cp, axis=1 ).get() - else: + except (ImportError, OSError): new_mesh.point_data["DeformationMagnitude"] = np.linalg.norm( np.asarray(new_pnts) - np.asarray(pnts), axis=1 ) diff --git a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py index bd58f0e..4ebd7cb 100644 --- a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py @@ -5,18 +5,18 @@ The workflow includes: 1. Rough alignment using ICP (RegisterModelsICP) 1.5. Optional PCA-based registration (RegisterModelsPCA) if PCA data provided -2. Mask-based deformable registration (RegisterModelsDistanceMaps) -3. Optional final mask-to-image refinement using Icon +2. Labelmap-based deformable registration (RegisterModelsDistanceMaps) +3. Optional final labelmap-to-image refinement using Icon The registration is particularly useful for cardiac modeling where a generic heart model needs to be fitted to patient-specific imaging data. Key Features: - - Automatic mask generation if not provided by user + - Automatic labelmap generation if not provided by user - Modular design using RegisterModelsICP, RegisterModelsPCA, and RegisterModelsDistanceMaps - Multi-stage registration pipeline: - ICP → (optional PCA) → mask-to-mask → mask-to-image + ICP → (optional PCA) → labelmap-to-labelmap → labelmap-to-image - Optional PCA-based shape fitting - Support for multi-label anatomical structures - Optional Icon-based final refinement @@ -41,57 +41,79 @@ from .workflow_convert_image_to_vtk import WorkflowConvertImageToVTK -def _load_tubetk() -> Any: - """Load TubeTK lazily for operations that require its filters.""" - from itk import TubeTK as ttk - - return ttk - - def _image_has_isotropic_spacing(image: itk.Image) -> bool: """Return whether all image spacing values match within numeric tolerance.""" spacing = np.asarray(image.GetSpacing(), dtype=np.float64) return bool(np.allclose(spacing, spacing[0])) +def _make_isotropic_image(image: itk.Image) -> itk.Image: + """Resample *image* to isotropic spacing using the finest voxel pitch. + + Equivalent to TubeTK's ResampleImage.SetMakeHighResIso(True), implemented + with standard ITK ResampleImageFilter so that TubeTK is not needed here. + """ + spacing = np.asarray(image.GetSpacing(), dtype=np.float64) + size = np.asarray(image.GetLargestPossibleRegion().GetSize(), dtype=np.int64) + + min_spacing = float(spacing.min()) + new_spacing = [min_spacing] * 3 + # Ceiling to avoid clipping the image boundary. + new_size = [int(np.ceil(size[i] * spacing[i] / min_spacing)) for i in range(3)] + + ImageType = type(image) + interpolator = itk.LinearInterpolateImageFunction[ImageType, itk.D].New() + resampler = itk.ResampleImageFilter[ImageType, ImageType].New() + resampler.SetInput(image) + resampler.SetInterpolator(interpolator) + resampler.SetOutputSpacing(new_spacing) + resampler.SetSize(new_size) + resampler.SetOutputOrigin(image.GetOrigin()) + resampler.SetOutputDirection(image.GetDirection()) + resampler.Update() + result = resampler.GetOutput() + result.DisconnectPipeline() + return result + + class WorkflowFitStatisticalModelToPatient(PhysioMotion4DBase): - """Register anatomical models using multi-stage ICP, mask-based, and image-based + """Register anatomical models using multi-stage ICP, labelmap-based, and image-based registration. This class provides a flexible workflow for registering generic anatomical models (e.g., cardiac models) to patient-specific surface models and images. The registration pipeline combines: - Initial model alignment using RegisterModelsICP (centroid + affine ICP) - - Mask-based deformable registration using RegisterModelsDistanceMaps (Greedy/ICON) - - Optional final mask-to-image refinement using Icon registration + - Labelmap-based deformable registration using RegisterModelsDistanceMaps (Greedy/ICON) + - Optional final labelmap-to-image refinement using Icon registration **Registration Pipeline:** 1. **ICP Alignment**: Rough affine alignment using RegisterModelsICP 2. **PCA Registration**: Performs PCA-based shape fitting using RegisterModelsPCA - 3. **Mask-to-Mask**: Deformable registration using RegisterModelsDistanceMaps - 4. **Mask-to-Image**: Final refinement + 3. **Labelmap-to-Labelmap**: Deformable registration using RegisterModelsDistanceMaps + 4. **Labelmap-to-Image**: Final refinement - **Mask Configuration:** - Masks are automatically generated from models if not provided by the user - via set_masks(). Auto-generated masks use mask_dilation_mm parameter. + **Labelmap Configuration:** + Labelmaps are automatically generated from models if not provided by the user + via set_labelmaps(). Auto-generated labelmaps use labelmap_dilation_mm parameter. Attributes: template_model (pv.DataSet): Generic anatomical model to be registered template_model_surface (pv.PolyData): Surface extracted from template_model_surface - template_model_mask (itk.Image): Binary/multi-label mask for model model - template_model_roi (itk.Image): ROI mask for model model + template_model_labelmap (itk.Image): Multi-label labelmap for template model + template_model_mask (itk.Image): Binary mask for template model registration region patient_models (list of pv.DataSet): Patient-specific models patient_model_surface (pv.PolyData): Primary patient model surface (first in list) combined_patient_model (pv.PolyData): Merged patient models before surface extraction; used when pca_uses_surface=False. patient_image (itk.Image): Reference image providing coordinate frame - patient_mask (itk.Image): Binary/multi-label mask for patient model - patient_roi (itk.Image): ROI mask for patient model - mask_dilation_mm (float): Dilation for mask generation - roi_dilation_mm (float): Dilation for ROI mask + patient_labelmap (itk.Image): Multi-label labelmap for patient model + patient_mask (itk.Image): Binary mask for patient registration region + labelmap_dilation_mm (float): Dilation for labelmap generation + mask_dilation_mm (float): Dilation for binary mask generation transform_tools (TransformTools): Transform utilities registrar_ICON (RegisterImagesICON): ICON registration instance registrar_Greedy (RegisterImagesGreedy): Greedy registration instance @@ -107,15 +129,15 @@ class WorkflowFitStatisticalModelToPatient(PhysioMotion4DBase): (if PCA used) pca_template_model_surface: template model surface after PCA registration (if PCA used) - m2m_forward_transform: Mask-to-mask forward transform - m2m_inverse_transform: Mask-to-mask inverse transform - m2m_template_model_surface: template model surface after mask-to-mask + l2l_forward_transform: Labelmap-to-labelmap forward transform + l2l_inverse_transform: Labelmap-to-labelmap inverse transform + l2l_template_model_surface: template model surface after labelmap-to-labelmap registration - m2i_forward_transform: Mask-to-image forward transform - m2i_inverse_transform: Mask-to-image inverse transform - m2i_template_model_surface: template model surface after mask-to-image + l2i_forward_transform: Labelmap-to-image forward transform + l2i_inverse_transform: Labelmap-to-image inverse transform + l2i_template_model_surface: template model surface after labelmap-to-image registration - m2i_template_labelmap: template labelmap after mask-to-image registration + l2i_template_labelmap: template labelmap after labelmap-to-image registration registered_template_model: Final registered model registered_template_model_surface: Final registered model surface @@ -125,11 +147,11 @@ class WorkflowFitStatisticalModelToPatient(PhysioMotion4DBase): ... template_model=heart_model, ... patient_models=[lv_model, mc_model, rv_model], ... ) - >>> registrar.set_roi_dilation_mm(20) + >>> registrar.set_mask_dilation_mm(20) >>> # To enable PCA registration, call before run_workflow(): >>> # registrar.set_use_pca_registration(True, pca_model=pca_model_dict, pca_number_of_modes=10) - >>> # To enable mask-to-image refinement: - >>> # registrar.set_use_mask_to_image_registration(True, template_labelmap, organ_mesh_ids, organ_extra_ids, background_ids) + >>> # To enable labelmap-to-image refinement: + >>> # registrar.set_use_labelmap_to_image_registration(True, template_labelmap, organ_mesh_ids, organ_extra_ids, background_ids) >>> result = registrar.run_workflow() """ @@ -205,11 +227,7 @@ def __init__( if patient_image is not None: self.patient_image = patient_image if not _image_has_isotropic_spacing(self.patient_image): - ttk = _load_tubetk() - resampler = ttk.ResampleImage.New(Input=self.patient_image) - resampler.SetMakeHighResIso(True) - resampler.Update() - self.patient_image = resampler.GetOutput() + self.patient_image = _make_isotropic_image(self.patient_image) else: self.patient_image = self.contour_tools.create_reference_image( mesh=self.patient_model_surface, @@ -220,22 +238,22 @@ def __init__( self.registrar_Greedy = RegisterImagesGreedy() self.registrar_Greedy.set_number_of_iterations([5, 2, 5]) - # Icon registration for final mask-to-image step + # Icon registration for final labelmap-to-image step self.registrar_ICON = RegisterImagesICON() self.registrar_ICON.set_modality("ct") self.registrar_ICON.set_mass_preservation(False) self.registrar_ICON.set_multi_modality(True) self.registrar_ICON.set_number_of_iterations(50) - # Mask configuration (auto-generated) + # Labelmap/mask configuration (auto-generated) + self.template_model_labelmap = None + self.patient_labelmap = None self.template_model_mask = None self.patient_mask = None - self.template_model_roi = None - self.patient_roi = None - # Parameters for mask generation and processing - self.mask_dilation_mm: float = 0.0 # For auto-generated mask dilation - self.roi_dilation_mm: float = 25.0 # For ROI mask generation + # Parameters for labelmap and mask generation + self.labelmap_dilation_mm: float = 0.0 # For auto-generated labelmap dilation + self.mask_dilation_mm: float = 25.0 # For binary registration mask generation # Stage 1: ICP alignment results self.icp_registrar: Optional[RegisterModelsICP] = None @@ -258,19 +276,19 @@ def __init__( self.pca_template_labelmap: Optional[itk.Image] = None self.pca_uses_surface: bool = False - # Stage 2: Mask-to-mask registration results - self.use_m2m_registration = True - self.m2m_inverse_transform: Optional[itk.Transform] = None - self.m2m_forward_transform: Optional[itk.Transform] = None - self.m2m_template_model_surface: Optional[pv.PolyData] = None - self.m2m_template_labelmap: Optional[itk.Image] = None + # Stage 2: Labelmap-to-labelmap registration results + self.use_l2l_registration = True + self.l2l_inverse_transform: Optional[itk.Transform] = None + self.l2l_forward_transform: Optional[itk.Transform] = None + self.l2l_template_model_surface: Optional[pv.PolyData] = None + self.l2l_template_labelmap: Optional[itk.Image] = None - # Stage 3: Mask-to-image registration results (disabled by default; enable via set_use_mask_to_image_registration(True, template_labelmap, ...)) - self.use_m2i_registration = False - self.m2i_inverse_transform: Optional[itk.Transform] = None - self.m2i_forward_transform: Optional[itk.Transform] = None - self.m2i_template_model_surface: Optional[pv.PolyData] = None - self.m2i_template_labelmap: Optional[itk.Image] = None + # Stage 3: Labelmap-to-image registration results (disabled by default; enable via set_use_labelmap_to_image_registration(True, template_labelmap, ...)) + self.use_l2i_registration = False + self.l2i_inverse_transform: Optional[itk.Transform] = None + self.l2i_forward_transform: Optional[itk.Transform] = None + self.l2i_template_model_surface: Optional[pv.PolyData] = None + self.l2i_template_labelmap: Optional[itk.Image] = None self.use_ICON_registration_refinement = False @@ -279,103 +297,24 @@ def __init__( self.registered_template_model_surface: Optional[pv.PolyData] = None self.registered_template_labelmap: Optional[itk.Image] = None - def _auto_generate_mask( - self, models: list[pv.DataSet], dilate_mm: Optional[float] = None - ) -> itk.Image: - """Auto-generate binary masks from models. - - Creates binary masks from list of models, with dilation - according to mask_dilation_mm parameter. - """ - self.log_info( - f"Auto-generating masks from models (dilation:{self.mask_dilation_mm}mm)..." - ) - - # Generate patient mask (single model or multi-label) - if len(models) == 1: - mask = self.contour_tools.create_mask_from_mesh( - models[0], - self.patient_image, - ) - else: - # Create multi-label mask - first_mask = self.contour_tools.create_mask_from_mesh( - models[0], - self.patient_image, - ) - mask_arr = np.zeros_like(itk.GetArrayFromImage(first_mask), dtype=np.uint8) - for i, model in enumerate(models): - if i == 0: - mask = first_mask - else: - mask = self.contour_tools.create_mask_from_mesh( - model, - self.patient_image, - ) - cur_mask = itk.GetArrayFromImage(mask) - mask_arr[cur_mask > 0] = i + 1 - mask = itk.GetImageFromArray(mask_arr.astype(np.uint8)) - mask.CopyInformation(self.patient_image) - - # Apply dilation if requested - if dilate_mm is None: - dilate_mm = self.mask_dilation_mm - if dilate_mm > 0: - mask = self.labelmap_tools.convert_labelmap_to_mask( - mask, dilation_in_mm=dilate_mm - ) - - self.log_info("Masks auto-generated successfully.") - - return mask - - def _auto_generate_roi_mask( - self, mask: itk.Image, dilate_mm: Optional[float] = None - ) -> itk.Image: - """Auto-generate ROI mask from existing masks with dilation. + def set_labelmap_dilation_mm(self, labelmap_dilation_mm: float) -> None: + """Set dilation amount for auto-generated labelmaps. - Uses self.roi_dilation_mm for dilation amount. - - Note: - Requires masks to exist (auto-generated or user-provided). + Args: + labelmap_dilation_mm: Dilation amount in millimeters applied when + auto-generating multi-label labelmaps from models. Default: 0mm """ - self.log_info( - f"Auto-generating ROI masks (dilation: {self.roi_dilation_mm}mm)..." - ) - - if dilate_mm is None: - dilate_mm = self.roi_dilation_mm - - # Generate model ROI mask - roi = None - if dilate_mm > 0: - roi = self.labelmap_tools.convert_labelmap_to_mask( - mask, dilation_in_mm=dilate_mm - ) - else: - roi = mask - - self.log_info("ROI masks auto-generated successfully.") - return roi + self.labelmap_dilation_mm = labelmap_dilation_mm def set_mask_dilation_mm(self, mask_dilation_mm: float) -> None: - """Set mask dilation amount for auto-generated masks. + """Set dilation amount for binary registration masks. Args: - mask_dilation_mm: Dilation amount in millimeters for mask generation. - Default: 5mm + mask_dilation_mm: Dilation amount in millimeters for binary registration + mask generation. Default: 25mm """ self.mask_dilation_mm = mask_dilation_mm - def set_roi_dilation_mm(self, roi_dilation_mm: float) -> None: - """Set ROI mask dilation amount. - - Args: - roi_dilation_mm: Dilation amount in millimeters for ROI mask generation. - Default: 20mm - """ - self.roi_dilation_mm = roi_dilation_mm - def set_use_pca_registration( self, use_pca_registration: bool, @@ -411,32 +350,32 @@ def set_use_pca_registration( self.pca_number_of_modes = 0 self.use_pca_registration = use_pca_registration - def set_use_mask_to_mask_registration( - self, use_mask_to_mask_registration: bool + def set_use_labelmap_to_labelmap_registration( + self, use_labelmap_to_labelmap_registration: bool ) -> None: - """Set whether to use mask-to-mask registration. + """Set whether to use labelmap-to-labelmap registration. Args: - use_mask_to_mask_registration: Whether to use mask-to-mask registration. - Default: True + use_labelmap_to_labelmap_registration: Whether to use labelmap-to-labelmap + deformable registration. Default: True """ - self.use_m2m_registration = use_mask_to_mask_registration + self.use_l2l_registration = use_labelmap_to_labelmap_registration - def set_use_mask_to_image_registration( + def set_use_labelmap_to_image_registration( self, - use_mask_to_image_registration: bool, + use_labelmap_to_image_registration: bool, template_labelmap: Optional[itk.Image] = None, template_labelmap_organ_mesh_ids: Optional[list[int]] = None, template_labelmap_organ_extra_ids: Optional[list[int]] = None, template_labelmap_background_ids: Optional[list[int]] = None, ) -> None: - """Set whether to use mask-to-image registration. + """Set whether to use labelmap-to-image registration. When enabling (True), a template labelmap and label IDs must be provided so the workflow can propagate and refine the labelmap to the patient image. Args: - use_mask_to_image_registration: Whether to use mask-to-image registration. + use_labelmap_to_image_registration: Whether to use labelmap-to-image registration. template_labelmap: Required when use is True. Template labelmap in template model space (same geometry as template_model). template_labelmap_organ_mesh_ids: Required when use is True. Label IDs for @@ -450,31 +389,31 @@ def set_use_mask_to_image_registration( ValueError: If use is True and any of template_labelmap or the id lists is None or missing. """ - if use_mask_to_image_registration: + if use_labelmap_to_image_registration: if template_labelmap is None: raise ValueError( - "When enabling mask-to-image registration, template_labelmap must be provided." + "When enabling labelmap-to-image registration, template_labelmap must be provided." ) if template_labelmap_organ_mesh_ids is None: raise ValueError( - "When enabling mask-to-image registration, " + "When enabling labelmap-to-image registration, " "template_labelmap_organ_mesh_ids must be provided." ) if template_labelmap_organ_extra_ids is None: raise ValueError( - "When enabling mask-to-image registration, " + "When enabling labelmap-to-image registration, " "template_labelmap_organ_extra_ids must be provided." ) if template_labelmap_background_ids is None: raise ValueError( - "When enabling mask-to-image registration, " + "When enabling labelmap-to-image registration, " "template_labelmap_background_ids must be provided." ) self.template_labelmap = template_labelmap self.template_labelmap_organ_mesh_ids = template_labelmap_organ_mesh_ids self.template_labelmap_organ_extra_ids = template_labelmap_organ_extra_ids self.template_labelmap_background_ids = template_labelmap_background_ids - self.use_m2i_registration = use_mask_to_image_registration + self.use_l2i_registration = use_labelmap_to_image_registration def _transform_model_dataset( self, @@ -680,69 +619,70 @@ def register_model_to_model_pca(self) -> dict: "registered_template_labelmap": self.pca_template_labelmap, } - def register_mask_to_mask(self) -> Optional[dict]: - """Perform mask-based deformable registration of model to patient model. + def register_labelmap_to_labelmap(self) -> Optional[dict]: + """Perform labelmap-based deformable registration of template model to patient model. - Uses RegisterModelsDistanceMaps with Greedy affine followed by ICON deformable registration. + Uses RegisterModelsDistanceMaps with Greedy affine followed by ICON deformable + registration on distance maps derived from the model surfaces. Returns: dict: Dictionary containing: - - 'forward_transform': model to patient space transform - - 'inverse_transform': patient to model space transform - - 'registered_template_model_surface': Transformed model model - - 'registered_template_labelmap': Transformed model labelmap + - 'forward_transform': template to patient space transform + - 'inverse_transform': patient to template space transform + - 'registered_template_model_surface': Transformed template model surface + - 'registered_template_labelmap': Transformed template labelmap """ self.log_section( - "Stage 3: Mask-to-Mask Deformable Registration", + "Stage 3: Labelmap-to-Labelmap Deformable Registration", width=70, ) - if not self.use_m2m_registration: - self.log_info("Mask-to-mask registration is not enabled.") + if not self.use_l2l_registration: + self.log_info("Labelmap-to-labelmap registration is not enabled.") return None - # Create mask-based registrar + # Create labelmap-based registrar assert self.pca_template_model_surface is not None, ( "PCA template model surface must be set" ) - mask_registrar = RegisterModelsDistanceMaps( + labelmap_registrar = RegisterModelsDistanceMaps( moving_model=self.pca_template_model_surface, fixed_model=self.patient_model_surface, reference_image=self.patient_image, - roi_dilation_mm=self.roi_dilation_mm, + mask_dilation_mm=self.mask_dilation_mm, ) # Run deformable registration - mask_result = mask_registrar.register( + l2l_result = labelmap_registrar.register( transform_type="Deformable", ) # Store results - self.m2m_forward_transform = mask_result["forward_transform"] - self.m2m_inverse_transform = mask_result["inverse_transform"] - self.m2m_template_model_surface = mask_result["registered_model"] + self.l2l_forward_transform = l2l_result["forward_transform"] + self.l2l_inverse_transform = l2l_result["inverse_transform"] + self.l2l_template_model_surface = l2l_result["registered_model"] - self.registered_template_model_surface = self.m2m_template_model_surface + self.registered_template_model_surface = self.l2l_template_model_surface if self.pca_template_labelmap is not None: - self.m2m_template_labelmap = self.transform_tools.transform_image( + self.l2l_template_labelmap = self.transform_tools.transform_image( self.pca_template_labelmap, - self.m2m_forward_transform, + self.l2l_forward_transform, self.patient_image, interpolation_method="nearest", ) else: - self.m2m_template_labelmap = None + self.l2l_template_labelmap = None - self.registered_template_labelmap = self.m2m_template_labelmap + self.registered_template_labelmap = self.l2l_template_labelmap - self.log_info("Stage 3 complete: Mask-to-mask registration finished.") + self.log_info("Stage 3 complete: Labelmap-to-labelmap registration finished.") return { - "forward_transform": self.m2m_forward_transform, - "inverse_transform": self.m2m_inverse_transform, - "registered_template_model_surface": self.m2m_template_model_surface, - "registered_template_labelmap": self.m2m_template_labelmap, + "forward_transform": self.l2l_forward_transform, + "inverse_transform": self.l2l_inverse_transform, + "registered_template_model_surface": self.l2l_template_model_surface, + "registered_template_labelmap": self.l2l_template_labelmap, } def register_labelmap_to_image( @@ -750,14 +690,15 @@ def register_labelmap_to_image( ) -> Optional[dict]: """Perform labelmap-to-image refinement. - Uses registration to align labelmap to actual image intensities. + Uses registration to align the propagated template labelmap to actual + image intensities. Returns: dict: Dictionary containing: - - 'inverse_transform': patient to model space transform - - 'forward_transform': model to patient space transform - - 'registered_template_model_surface': Transformed model model - - 'registered_template_labelmap': Transformed model labelmap + - 'inverse_transform': patient to template space transform + - 'forward_transform': template to patient space transform + - 'registered_template_model_surface': Transformed template model surface + - 'registered_template_labelmap': Transformed template labelmap """ self.log_section( "Stage 4: Labelmap-to-Image Refinement (Icon Registration)", width=70 @@ -770,99 +711,101 @@ def register_labelmap_to_image( or self.template_labelmap_background_ids is None ): raise ValueError( - "Mask-to-image registration requires template labelmap and label IDs. " - "Call set_use_mask_to_image_registration(True, template_labelmap, " + "Labelmap-to-image registration requires template labelmap and label IDs. " + "Call set_use_labelmap_to_image_registration(True, template_labelmap, " "organ_mesh_ids, organ_extra_ids, background_ids) before run_workflow()." ) - if self.m2m_template_labelmap is None: + if self.l2l_template_labelmap is None: raise ValueError( - "Mask-to-image registration requires a labelmap to have been set " - "(via set_use_mask_to_image_registration(True, ...)) before running " - "earlier stages so the labelmap is propagated through ICP/PCA/M2M." + "Labelmap-to-image registration requires a labelmap to have been set " + "(via set_use_labelmap_to_image_registration(True, ...)) before running " + "earlier stages so the labelmap is propagated through ICP/PCA/L2L." ) - labelmap_arr = itk.GetArrayFromImage(self.m2m_template_labelmap).astype( - np.uint16 - ) - labelmap_arr = np.where( - np.isin(labelmap_arr, self.template_labelmap_background_ids), + template_labelmap_arr = itk.GetArrayFromImage( + self.l2l_template_labelmap + ).astype(np.uint16) + template_labelmap_arr = np.where( + np.isin(template_labelmap_arr, self.template_labelmap_background_ids), 0, - labelmap_arr, + template_labelmap_arr, ) - labelmap_arr = np.where( - np.isin(labelmap_arr, self.template_labelmap_organ_mesh_ids), + template_labelmap_arr = np.where( + np.isin(template_labelmap_arr, self.template_labelmap_organ_mesh_ids), 0, - labelmap_arr, + template_labelmap_arr, ) - labelmap_arr = np.where( - np.isin(labelmap_arr, self.template_labelmap_organ_extra_ids), + template_labelmap_arr = np.where( + np.isin(template_labelmap_arr, self.template_labelmap_organ_extra_ids), 1, - labelmap_arr, + template_labelmap_arr, ) - labelmap = itk.GetImageFromArray(labelmap_arr) - labelmap.CopyInformation(self.m2m_template_labelmap) + template_labelmap = itk.GetImageFromArray(template_labelmap_arr) + template_labelmap.CopyInformation(self.l2l_template_labelmap) - labelmap_roi = self._auto_generate_roi_mask(labelmap) + template_labelmap_mask = self.labelmap_tools.convert_labelmap_to_mask( + template_labelmap, dilation_in_mm=self.mask_dilation_mm + ) - patient_mask = self._auto_generate_mask( - [self.patient_model_surface], dilate_mm=0 + patient_mask = self.contour_tools.create_mask_from_mesh( + self.patient_model_surface, + self.patient_image, ) - patient_roi = self._auto_generate_roi_mask(patient_mask) self.registrar_Greedy.set_fixed_image(self.patient_image) - self.registrar_Greedy.set_fixed_mask(patient_roi) + self.registrar_Greedy.set_fixed_mask(patient_mask) result = self.registrar_Greedy.register( - moving_image=labelmap, moving_mask=labelmap_roi + moving_image=template_labelmap, moving_mask=template_labelmap_mask ) - self.m2i_inverse_transform = result["inverse_transform"] - self.m2i_forward_transform = result["forward_transform"] + self.l2i_inverse_transform = result["inverse_transform"] + self.l2i_forward_transform = result["forward_transform"] if use_ICON_refinement: # Configure Icon registration self.registrar_ICON.set_fixed_image(self.patient_image) - self.registrar_ICON.set_fixed_mask(patient_roi) + self.registrar_ICON.set_fixed_mask(patient_mask) # Perform Icon registration result = self.registrar_ICON.register( - initial_forward_transform=self.m2i_forward_transform, - moving_image=labelmap, - moving_mask=labelmap_roi, + initial_forward_transform=self.l2i_forward_transform, + moving_image=template_labelmap, + moving_mask=template_labelmap_mask, ) - self.m2i_inverse_transform = result["inverse_transform"] - self.m2i_forward_transform = result["forward_transform"] + self.l2i_inverse_transform = result["inverse_transform"] + self.l2i_forward_transform = result["forward_transform"] - # Transform model with Icon result - assert self.m2m_template_model_surface is not None, ( - "M2M template model surface must be set" + # Transform model with result + assert self.l2l_template_model_surface is not None, ( + "L2L template model surface must be set" ) - self.m2i_template_model_surface = cast( + self.l2i_template_model_surface = cast( pv.PolyData, self._transform_model_dataset( - self.m2m_template_model_surface, - self.m2i_inverse_transform, + self.l2l_template_model_surface, + self.l2i_inverse_transform, with_deformation_magnitude=True, ), ) - self.m2i_template_labelmap = self.transform_tools.transform_image( + self.l2i_template_labelmap = self.transform_tools.transform_image( self.template_labelmap, - self.m2i_forward_transform, + self.l2i_forward_transform, self.patient_image, interpolation_method="nearest", ) - self.log_info("Stage 4 complete: Mask-to-image registration finished.") + self.log_info("Stage 4 complete: Labelmap-to-image registration finished.") - self.registered_template_model_surface = self.m2i_template_model_surface + self.registered_template_model_surface = self.l2i_template_model_surface - self.registered_template_labelmap = self.m2i_template_labelmap + self.registered_template_labelmap = self.l2i_template_labelmap return { - "inverse_transform": self.m2i_inverse_transform, - "forward_transform": self.m2i_forward_transform, - "registered_template_model_surface": self.m2i_template_model_surface, - "registered_template_labelmap": self.m2i_template_labelmap, + "inverse_transform": self.l2i_inverse_transform, + "forward_transform": self.l2i_forward_transform, + "registered_template_model_surface": self.l2i_template_model_surface, + "registered_template_labelmap": self.l2i_template_labelmap, } def transform_model( @@ -905,10 +848,10 @@ def transform_model( transform_steps.append( ("PCA post-transform", self.pca_registrar.post_pca_transform) ) - if self.use_m2m_registration and self.m2m_inverse_transform is not None: - transform_steps.append(("Mask-to-mask", self.m2m_inverse_transform)) - if self.use_m2i_registration and self.m2i_inverse_transform is not None: - transform_steps.append(("Mask-to-image", self.m2i_inverse_transform)) + if self.use_l2l_registration and self.l2l_inverse_transform is not None: + transform_steps.append(("Labelmap-to-labelmap", self.l2l_inverse_transform)) + if self.use_l2i_registration and self.l2i_inverse_transform is not None: + transform_steps.append(("Labelmap-to-image", self.l2i_inverse_transform)) for i, (name, tfm) in enumerate(transform_steps, start=1): self.log_progress(i, len(transform_steps), prefix=f"Applying {name}") @@ -940,14 +883,14 @@ def run_workflow( 1. ICP alignment (RegisterModelsICP) 2. PCA registration (PCA data was provided) - 3. Mask-to-mask deformable registration (RegisterModelsDistanceMaps) - 4. Optional mask-to-image refinement (Icon); requires template labelmap and IDs - set via set_use_mask_to_image_registration(True, ...). + 3. Labelmap-to-labelmap deformable registration (RegisterModelsDistanceMaps) + 4. Optional labelmap-to-image refinement (Icon); requires template labelmap and IDs + set via set_use_labelmap_to_image_registration(True, ...). Args: use_ICON_registration_refinement: Whether to apply ICON refinement in the - mask-to-image stage (Stage 4). The mask-to-mask stage always uses - Greedy affine + ICON deformable. Default: False + labelmap-to-image stage (Stage 4). The labelmap-to-labelmap stage always + uses Greedy affine + ICON deformable. Default: False Returns: dict with registered_template_model and registered_template_model_surface @@ -964,12 +907,12 @@ def run_workflow( # Stage 2: Optional PCA registration (if PCA data was set) self.register_model_to_model_pca() - # Stage 3: Optional Mask-to-mask deformable registration - if self.use_m2m_registration: - self.register_mask_to_mask() + # Stage 3: Optional Labelmap-to-labelmap deformable registration + if self.use_l2l_registration: + self.register_labelmap_to_labelmap() - # Stage 4: Optional mask-to-image refinement - if self.use_m2i_registration: + # Stage 4: Optional labelmap-to-image refinement + if self.use_l2i_registration: self.register_labelmap_to_image( use_ICON_refinement=use_ICON_registration_refinement ) diff --git a/tests/test_workflow_fit_statistical_model_to_patient.py b/tests/test_workflow_fit_statistical_model_to_patient.py index 0444203..f8165e5 100644 --- a/tests/test_workflow_fit_statistical_model_to_patient.py +++ b/tests/test_workflow_fit_statistical_model_to_patient.py @@ -16,10 +16,10 @@ ) -def test_auto_generate_mask_accumulates_multilabel_models( +def test_auto_generate_labelmap_accumulates_multilabel_models( monkeypatch: Any, ) -> None: - """Multi-model masks accumulate label IDs instead of overwriting prior labels.""" + """Multi-model labelmaps accumulate label IDs instead of overwriting prior labels.""" image = itk.image_from_array(np.zeros((3, 3, 3), dtype=np.float32)) model = pv.PolyData( np.array( @@ -53,7 +53,7 @@ def fake_create_mask_from_mesh(*_args: object, **_kwargs: object) -> Any: fake_create_mask_from_mesh, ) - output = workflow._auto_generate_mask([model, model], dilate_mm=0.0) + output = workflow._auto_generate_labelmap([model, model], dilate_mm=0.0) output_arr = itk.GetArrayFromImage(output) assert output_arr[0, 0, 0] == 1 diff --git a/tutorials/tutorial_07_dirlab_pca_model.py b/tutorials/tutorial_07_dirlab_pca_model.py index b0dc589..fdd746a 100644 --- a/tutorials/tutorial_07_dirlab_pca_model.py +++ b/tutorials/tutorial_07_dirlab_pca_model.py @@ -260,7 +260,7 @@ def fit_model( pca_number_of_modes=0, pca_uses_surface=True, ) - workflow.set_use_mask_to_mask_registration(False) + workflow.set_use_labelmap_to_labelmap_registration(False) result = workflow.run_workflow() fitted_surface = result["registered_template_model_surface"] From ac9c44c6d826c7631bfb5c78237a37129f540cfc Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Wed, 17 Jun 2026 09:14:11 -0400 Subject: [PATCH 2/6] BUG: Remove test of removed function --- ...rkflow_fit_statistical_model_to_patient.py | 45 ------------------- 1 file changed, 45 deletions(-) diff --git a/tests/test_workflow_fit_statistical_model_to_patient.py b/tests/test_workflow_fit_statistical_model_to_patient.py index f8165e5..f543cbc 100644 --- a/tests/test_workflow_fit_statistical_model_to_patient.py +++ b/tests/test_workflow_fit_statistical_model_to_patient.py @@ -16,51 +16,6 @@ ) -def test_auto_generate_labelmap_accumulates_multilabel_models( - monkeypatch: Any, -) -> None: - """Multi-model labelmaps accumulate label IDs instead of overwriting prior labels.""" - image = itk.image_from_array(np.zeros((3, 3, 3), dtype=np.float32)) - model = pv.PolyData( - np.array( - [ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 1.0, 0.0], - ], - dtype=np.float64, - ) - ) - workflow = WorkflowFitStatisticalModelToPatient( - template_model=model, - patient_models=[model], - patient_image=image, - ) - workflow.patient_image = image - masks = [ - np.pad(np.ones((1, 1, 1), dtype=np.uint8), ((0, 2), (0, 2), (0, 2))), - np.pad(np.ones((1, 1, 1), dtype=np.uint8), ((2, 0), (2, 0), (2, 0))), - ] - - def fake_create_mask_from_mesh(*_args: object, **_kwargs: object) -> Any: - mask = itk.image_from_array(masks.pop(0)) - mask.CopyInformation(workflow.patient_image) - return mask - - monkeypatch.setattr( - workflow.contour_tools, - "create_mask_from_mesh", - fake_create_mask_from_mesh, - ) - - output = workflow._auto_generate_labelmap([model, model], dilate_mm=0.0) - output_arr = itk.GetArrayFromImage(output) - - assert output_arr[0, 0, 0] == 1 - assert output_arr[2, 2, 2] == 2 - assert sorted(np.unique(output_arr).tolist()) == [0, 1, 2] - - def test_transform_model_applies_staged_transform() -> None: """Transform helper updates mesh points with image shape (Z, Y, X) = (3, 3, 3).""" image = itk.image_from_array(np.zeros((3, 3, 3), dtype=np.float32)) From 16e089c7d19c1aefb5ec336d93fb67f00fc2a094 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sat, 20 Jun 2026 11:57:36 -0400 Subject: [PATCH 3/6] BUG: Suggestions from coderabbit --- .../heart_model_to_patient-CHOPValve.py | 14 +++++--------- .../cli/fit_statistical_model_to_patient.py | 18 ++++++++++++------ src/physiomotion4d/contour_tools.py | 2 -- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py index 0607974..883b735 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py @@ -21,21 +21,21 @@ # %% # Patient CT image (defines coordinate frame) -patient_data_dir = Path.cwd().parent.parent / "data" / "CHOP-Valve4D" / "CT" +patient_data_dir = Path(__file__).parent.parent.parent / "data" / "CHOP-Valve4D" / "CT" patient_ct_path = patient_data_dir / "RVOT28-Dias.mha" # Template model (moving) -model_data_dir = Path.cwd().parent.parent / "data" / "KCL-Heart-Model" +model_data_dir = Path(__file__).parent.parent.parent / "data" / "KCL-Heart-Model" model_labelmap_path = model_data_dir / "labelmap" / "average_labelmap_with_bkg.mha" model_pca_data_dir = ( - Path.cwd().parent / "Heart-Create_Statistical_Model" / "kcl-heart-model" + Path(__file__).parent.parent / "Heart-Create_Statistical_Model" / "kcl-heart-model" ) model_pca_json_path = model_pca_data_dir / "pca_model.json" model_mesh_path = model_pca_data_dir / "pca_mean.vtp" model_pca_n_modes = 10 # Output directory -output_dir = Path.cwd() / "results-chop" +output_dir = Path(__file__).parent / "results-chop" # %% patient_image = itk.imread(str(patient_ct_path)) @@ -58,7 +58,7 @@ True, pca_model=model_pca_data, pca_number_of_modes=model_pca_n_modes ) -registrar.set_use_labelmap_to_labelmap_registration(True) +# registrar.set_use_labelmap_to_labelmap_registration(True) # %% patient_image = registrar.patient_image @@ -75,10 +75,6 @@ registered_model.save(str(output_dir / "registered_model.vtp")) registered_model_surface.save(str(output_dir / "registered_model_surface.vtp")) -registered_labelmap = results["registered_template_labelmap"] -itk.imwrite( - registered_labelmap, str(output_dir / "registered_labelmap.mha"), compression=True -) # %% pca_model = registrar.pca_template_model diff --git a/src/physiomotion4d/cli/fit_statistical_model_to_patient.py b/src/physiomotion4d/cli/fit_statistical_model_to_patient.py index 74b3a05..1bdd7f0 100644 --- a/src/physiomotion4d/cli/fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/cli/fit_statistical_model_to_patient.py @@ -240,10 +240,11 @@ def main() -> int: pca_model=pca_model, pca_number_of_modes=args.pca_number_of_modes, ) - if args.use_labelmap_to_labelmap: - workflow.set_use_labelmap_to_labelmap_registration( - args.use_labelmap_to_labelmap - ) + + workflow.set_use_labelmap_to_labelmap_registration( + args.use_labelmap_to_labelmap + ) + if args.use_labelmap_to_image: workflow.set_use_labelmap_to_image_registration( True, @@ -252,6 +253,7 @@ def main() -> int: template_labelmap_organ_extra_ids=args.template_labelmap_chamber_ids, template_labelmap_background_ids=args.template_labelmap_background_ids, ) + except (ValueError, RuntimeError, OSError) as e: print(f"Error initializing workflow: {e}") traceback.print_exc() @@ -290,13 +292,17 @@ def main() -> int: output_labelmap_file = os.path.join( args.output_dir, f"{args.output_prefix}_labelmap.nii.gz" ) - itk.imwrite(workflow.l2i_template_labelmap, output_labelmap_file) + itk.imwrite( + workflow.l2i_template_labelmap, output_labelmap_file, compression=True + ) print(f" Registered labelmap: {output_labelmap_file}") elif workflow.l2l_template_labelmap is not None: output_labelmap_file = os.path.join( args.output_dir, f"{args.output_prefix}_labelmap.nii.gz" ) - itk.imwrite(workflow.l2l_template_labelmap, output_labelmap_file) + itk.imwrite( + workflow.l2l_template_labelmap, output_labelmap_file, compression=True + ) print(f" Registered labelmap: {output_labelmap_file}") # Save intermediate results if available diff --git a/src/physiomotion4d/contour_tools.py b/src/physiomotion4d/contour_tools.py index b2ab147..9693a6a 100644 --- a/src/physiomotion4d/contour_tools.py +++ b/src/physiomotion4d/contour_tools.py @@ -39,8 +39,6 @@ def extract_contours( Args: labelmap_image (itk.image): The labelmap image to create contours from - output_file (str, optional): If provided, save the contours to this VTP - file Returns: pv.PolyData: The contours as a PyVista PolyData object From 75245ffe78545c6dc61e3d7d7e5e8a2c4c9fcc70 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sat, 20 Jun 2026 12:06:21 -0400 Subject: [PATCH 4/6] BUG: CodeRabbit --- src/physiomotion4d/labelmap_tools.py | 43 +++++++++++++++++++--------- 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/src/physiomotion4d/labelmap_tools.py b/src/physiomotion4d/labelmap_tools.py index 3baf22d..737b7b4 100644 --- a/src/physiomotion4d/labelmap_tools.py +++ b/src/physiomotion4d/labelmap_tools.py @@ -99,6 +99,9 @@ def create_distance_map( labelmap: itk.Image, max_distance_mm: float = 20.0, distance_scale: float = 5.0, + preserve_labels: bool = True, + fill_background_only: bool = False, + exclude_labels: Optional[list[int]] = None, ) -> itk.Image: """Encode a labelmap as a continuous label-plus-boundary-distance image. @@ -140,17 +143,24 @@ def create_distance_map( """ labels = itk.array_from_image(labelmap) - # A voxel is on a label boundary when it differs from a 6-connected - # neighbor along any axis. Mark both voxels straddling each change. - boundary = np.zeros(labels.shape, dtype=bool) - for axis in range(labels.ndim): - changed = np.diff(labels, axis=axis) != 0 - lower = [slice(None)] * labels.ndim - upper = [slice(None)] * labels.ndim - lower[axis] = slice(0, -1) - upper[axis] = slice(1, None) - boundary[tuple(lower)] |= changed - boundary[tuple(upper)] |= changed + if exclude_labels: + labels = np.where(np.isin(labels, exclude_labels), 0, labels) + + if not fill_background_only: + # A voxel is on a label boundary when it differs from a 6-connected + # neighbor along any axis. Mark both voxels straddling each change. + boundary = np.zeros(labels.shape, dtype=bool) + for axis in range(labels.ndim): + changed = np.diff(labels, axis=axis) != 0 + lower = [slice(None)] * labels.ndim + upper = [slice(None)] * labels.ndim + lower[axis] = slice(0, -1) + upper[axis] = slice(1, None) + boundary[tuple(lower)] |= changed + boundary[tuple(upper)] |= changed + else: + boundary = np.zeros(labels.shape, dtype=np.float32) + boundary[labels > 0] = 1.0 if boundary.any(): boundary_image = itk.image_from_array(boundary.astype(np.uint8)) @@ -161,16 +171,21 @@ def create_distance_map( distance_filter.SetSquaredDistance(False) distance_filter.SetUseImageSpacing(True) distance_filter.Update() - distance = np.abs( - itk.array_from_image(distance_filter.GetOutput()).astype(np.float32) + distance = itk.array_from_image(distance_filter.GetOutput()).astype( + np.float32 ) + if not fill_background_only: + distance = np.abs(distance) else: # No inter-label boundary exists (single uniform label); every # voxel gets a zero offset. distance = np.zeros(labels.shape, dtype=np.float32) offset = np.clip(distance, 0.0, max_distance_mm) / distance_scale - encoded = labels.astype(np.float32) + offset + if preserve_labels: + encoded = labels.astype(np.float32) + offset + else: + encoded = offset encoded_image = itk.image_from_array(encoded) encoded_image.CopyInformation(labelmap) From e8f2ca35e88fcda612681d80fd222a6f145e3715 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sat, 20 Jun 2026 12:13:24 -0400 Subject: [PATCH 5/6] ENH: Update env for cupy commit hooks --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 3c45de0..e43af6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -221,6 +221,8 @@ show_error_codes = true module = [ "ants", "cupy", + "cupy_backends", + "cupy_backends.*", "icon_registration", "icon_registration.*", "itk", From 4008048d98fb6097444e24e189a34f4c33e2ef02 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sat, 20 Jun 2026 12:14:44 -0400 Subject: [PATCH 6/6] BUG: Force commit of cupy methods until commit hooks updated --- src/physiomotion4d/transform_tools.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/src/physiomotion4d/transform_tools.py b/src/physiomotion4d/transform_tools.py index 4838fbc..f6f5eb5 100644 --- a/src/physiomotion4d/transform_tools.py +++ b/src/physiomotion4d/transform_tools.py @@ -14,7 +14,7 @@ import logging from collections.abc import Sequence from pathlib import Path -from typing import TypeAlias, cast +from typing import Type, TypeAlias, cast import itk import numpy as np @@ -364,13 +364,24 @@ def transform_dataset( if with_deformation_magnitude: try: import cupy as cp # noqa: PLC0415 - - new_pnts_cp = cp.array(new_pnts) - pnts_cp = cp.array(pnts) - new_mesh.point_data["DeformationMagnitude"] = cp.linalg.norm( - new_pnts_cp - pnts_cp, axis=1 - ).get() - except (ImportError, OSError): + except ImportError: + cp = None + if cp is not None: + try: + import cupy_backends.cuda.api.runtime as _cuda_rt # noqa: PLC0415 + + _CUDARuntimeError: Type[BaseException] = _cuda_rt.CUDARuntimeError + except ImportError: + _CUDARuntimeError = OSError + try: + new_pnts_cp = cp.array(new_pnts) + pnts_cp = cp.array(pnts) + new_mesh.point_data["DeformationMagnitude"] = cp.linalg.norm( + new_pnts_cp - pnts_cp, axis=1 + ).get() + except (OSError, _CUDARuntimeError): + cp = None + if cp is None: new_mesh.point_data["DeformationMagnitude"] = np.linalg.norm( np.asarray(new_pnts) - np.asarray(pnts), axis=1 )