From 728149805117acc324e652c8c092935158f26457 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 06:31:26 -0400 Subject: [PATCH 1/5] ENH: Rewrite PCA model registration around an analytic-gradient objective RegisterModelsPCA previously maximized mean intensity sampled from the fixed distance map through ITK's LinearInterpolateImageFunction, with the optimizer estimating gradients by finite differences. The objective is now stated and minimized directly: mean distance(model -> target) + w * mean distance(target -> model) # symmetric term + lambda * sum(b_i^2) # Mahalanobis shape prior Because the PCA deformation is linear in the coefficients b, the gradient is analytic and handed to the optimizer instead of being estimated, which removes one objective evaluation per coefficient per step. The post-PCA transform is folded into the mode directions so the gradient stays exact; when that transform is not affine its Jacobian is not constant, so the analytic gradient is disabled and finite differences are used with a logged warning. - Add symmetric_weight (default 0.5) so partial target coverage is penalized, and pca_prior_weight (default 0.0, disabled) for the Mahalanobis prior - Sample the distance map and its gradient with scipy.ndimage.map_coordinates, and build the target-to-model term with a scipy.spatial.cKDTree - Replace the cached ITK interpolator and _create_itk_points with _prepare_sampling, which builds the arrays the objective is made of once - Add ContourTools.sample_mesh_faces for face-density-aware point sampling and a negative_inside option on the signed distance map - Log transform fidelity after computing the PCA transforms Tests cover the pieces that were previously unverified: the analytic gradient against finite differences, recovery of known coefficients, the symmetric term penalizing partial coverage, the prior shrinking coefficients, eigenvector scaling by standard deviation, deformation happening in the template frame, and a transform round trip. --- src/physiotwin4d/contour_tools.py | 130 +++- src/physiotwin4d/register_models_pca.py | 777 ++++++++++++++++-------- tests/test_register_models_pca.py | 273 ++++++++- 3 files changed, 876 insertions(+), 304 deletions(-) diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 67558d1f..3d7d2c73 100644 --- a/src/physiotwin4d/contour_tools.py +++ b/src/physiotwin4d/contour_tools.py @@ -318,6 +318,54 @@ def create_labelmap_from_meshes( return labelmap_image + @staticmethod + def sample_mesh_faces(mesh: pv.DataSet, max_spacing: float) -> np.ndarray: + """Return mesh points supplemented by samples across the triangle faces. + + Rasterizing vertices alone leaves gaps between them on meshes that are + coarse relative to the voxel size, which makes a distance map built from + them ripple. Adding barycentric samples dense enough that consecutive + samples are closer than ``max_spacing`` closes those gaps. + + Args: + mesh: Source mesh; its surface is triangulated if needed. + max_spacing: Target spacing between samples, in mm. + + Returns: + (n, 3) array of sample points, starting with the mesh's own points. + """ + points = np.asarray(mesh.points, dtype=np.float64) + surface = mesh.extract_surface() if not isinstance(mesh, pv.PolyData) else mesh + surface = surface.triangulate() + if surface.faces.size == 0: + return points + faces = surface.faces.reshape(-1, 4)[:, 1:] + + vertices = np.asarray(surface.points, dtype=np.float64) + corners = vertices[faces] # (n_faces, 3, 3) + edge_lengths = np.linalg.norm( + corners - np.roll(corners, 1, axis=1), axis=2 + ).max(axis=1) + + # Group faces by how finely they need to be subdivided so each division + # level is generated as one vectorized batch. + divisions = np.maximum(1, np.ceil(edge_lengths / max(max_spacing, 1e-6))) + divisions = np.minimum(divisions, 64).astype(np.int64) + + samples = [points] + for level in np.unique(divisions): + if level < 2: + continue + selected = corners[divisions == level] + # Barycentric lattice with `level` divisions per edge. + steps = np.arange(level + 1, dtype=np.float64) / level + u, v = np.meshgrid(steps, steps, indexing="ij") + mask = (u + v) <= 1.0 + weights = np.column_stack([1.0 - u[mask] - v[mask], u[mask], v[mask]]) + samples.append(np.einsum("fca,kc->fka", selected, weights).reshape(-1, 3)) + + return np.concatenate(samples, axis=0) + def create_distance_map( self, mesh: pv.DataSet | pv.UnstructuredGrid, @@ -326,37 +374,62 @@ def create_distance_map( negative_inside: bool = True, zero_inside: bool = False, norm_to_max_distance: float = 0.0, + sample_faces: bool = True, ) -> itk.Image: - self.log_info("Computing signed distance map...") + """Compute a distance map of a mesh on the reference image's grid. - # Convert mask to binary - points = mesh.points + Args: + mesh: Mesh whose surface the distances are measured to. + reference_image: Image defining the output grid. + squared_distance: Sign-preserving square of the result. Default: False + negative_inside: Keep the signed output. Default: True + zero_inside: Clip negative values to zero before anything else. + Default: False + norm_to_max_distance: If non-zero, divide by this value and clip to + [-1, 1]. Default: 0.0 (distances stay in mm) + sample_faces: Rasterize samples across the triangle faces as well as + the vertices, so that coarse meshes do not leave gaps in the + rasterized surface. Default: True + + Returns: + ITK image of distances on the reference grid. + """ + self.log_info("Computing signed distance map...") size = reference_image.GetLargestPossibleRegion().GetSize() + if sample_faces: + points = self.sample_mesh_faces( + mesh, 0.5 * float(min(reference_image.GetSpacing())) + ) + self.log_debug( + "Distance map: %d face samples from %d mesh points", + len(points), + mesh.n_points, + ) + else: + points = np.asarray(mesh.points, dtype=np.float64) + # 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: - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - indx = reference_image.TransformPhysicalPointToIndex(itk_point) - if ( - indx[0] < 0 - or indx[1] < 0 - or indx[2] < 0 - or indx[0] >= size[0] - or indx[1] >= size[1] - or indx[2] >= size[2] - ): - continue - tmp_arr[indx[2], indx[1], indx[0]] = 1 - point_count += 1 + tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.uint8) + + # Bulk equivalent of TransformPhysicalPointToIndex, which rounds half up. + index_to_world = itk.array_from_matrix( + reference_image.GetDirection() + ) @ np.diag(np.asarray(reference_image.GetSpacing())) + origin = np.asarray(reference_image.GetOrigin(), dtype=np.float64) + indices = np.floor( + (points - origin) @ np.linalg.inv(index_to_world).T + 0.5 + ).astype(np.int64) + size_arr = np.array([size[0], size[1], size[2]], dtype=np.int64) + inside = np.all((indices >= 0) & (indices < size_arr), axis=1) + indices = indices[inside] + point_count = len(indices) + if point_count: + tmp_arr[indices[:, 2], indices[:, 1], indices[:, 0]] = 1 self.log_info( - "Distance map: %d/%d surface points within reference image", + "Distance map: %d/%d surface samples within reference image", point_count, len(points), ) @@ -370,8 +443,17 @@ def create_distance_map( str(size), str(reference_image.GetSpacing()), ) + elif not inside.all(): + # Distances near the dropped region are measured to whatever samples + # remain in the grid, so they are larger than the true distance. + self.log_warning( + "%d of %d surface samples fall outside the reference image; " + "distances near that boundary are overestimated.", + len(points) - point_count, + len(points), + ) - tmp_binary_image = itk.GetImageFromArray(tmp_arr.astype(np.uint8)) + tmp_binary_image = itk.GetImageFromArray(tmp_arr) tmp_binary_image.CopyInformation(reference_image) assert ( tmp_binary_image.GetLargestPossibleRegion().GetSize() diff --git a/src/physiotwin4d/register_models_pca.py b/src/physiotwin4d/register_models_pca.py index ef93d97c..c2cf68d0 100644 --- a/src/physiotwin4d/register_models_pca.py +++ b/src/physiotwin4d/register_models_pca.py @@ -8,7 +8,9 @@ import itk import numpy as np import pyvista as pv +from scipy.ndimage import map_coordinates from scipy.optimize import minimize +from scipy.spatial import cKDTree from typing_extensions import Self from .contour_tools import ContourTools @@ -17,29 +19,44 @@ class RegisterModelsPCA(PhysioTwin4DBase): - """Register PCA-based shape models to medical images using mean distance optimization. + """Register PCA-based shape models to images by minimizing a distance metric. This class implements a registration pipeline for fitting statistical shape models to patient-specific medical images: **PCA Deformable Registration** - Optimizes PCA coefficients - - Model equation: P = mean + Σ(b_i * std_i * pca_eigenvector_i) - - Maximizes mean distance at deformed model points P + - Model equation: P = template + Σ(b_i * std_i * pca_eigenvector_i) + - Minimizes the mean distance-to-target at the deformed model points P **Optimization Objective:** - Maximize the mean distance of the image sampled at model points using - ITK's LinearInterpolateImageFunction. This aligns the model with bright - regions in contrast-enhanced images (e.g., blood pool in cardiac CT). + ``fixed_distance_map`` is zero on the target surface and grows with + distance away from it (in mm), so the objective is *minimized*:: + + E(b) = (1 - w) * mean_i D(P_i(b)) # model -> target + + w * mean_j ||Q_j - P_nn(j)|| # target -> model + + lambda * Σ b_i² # Mahalanobis prior + + ``w`` is ``symmetric_weight`` and ``lambda`` is ``pca_prior_weight``. + Because the deformation is linear in ``b``, the gradient is analytic and + is supplied to the optimizer directly. + + **Coordinate frames:** + The eigenvectors are directions in the statistical model's own training + frame, so they are only valid when added to a template in that same + frame. Any rigid/affine alignment to the target must be supplied as + ``post_pca_transform``, which is applied *after* the deformation, rather + than pre-applied to ``pca_template_model``. Attributes: pca_template_model (pv.DataSet): Mean shape model pca_eigenvectors (np.ndarray): PCA eigenvectors/components (modes × n_points*3) pca_std_deviations (np.ndarray): Standard deviations per mode (modes,) - fixed_distance_map (itk.Image): Patient image providing distance data - n_points (int): Number of points in the model + fixed_distance_map (itk.Image): Distance map of the target, in mm + fixed_model (pv.DataSet): Target model, when one was supplied. Required + for the symmetric (target-to-model) term. pca_number_of_modes (int): Number of PCA modes available - pca_coefficients (np.ndarray): Optimized PCA coefficients + registered_model_pca_coefficients (np.ndarray): Optimized PCA coefficients registered_model (pv.DataSet): Final registered and deformed model post_pca_transform (itk.Transform): Transform to apply after PCA registration forward_point_transform (itk.DisplacementFieldTransform): POINT transform @@ -90,63 +107,99 @@ def __init__( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ): """Initialize the PCA-based model-to-image registration. Args: pca_template_model: PyVista model containing the mean 3D shape model - (unstructured grid or polydata) + (unstructured grid or polydata). It must be in the same frame + the PCA modes were trained in; supply any alignment to the + target as post_pca_transform instead of pre-applying it here. pca_eigenvectors: Numpy array of PCA eigenvectors/components. Shape: (modes, n_points*3) Each row is a flattened eigenmode with 3D displacements: [x1,y1,z1, x2,y2,z2, ...] pca_std_deviations: Numpy array of standard deviations per PCA mode. Shape: (modes,) These are the square roots of pca_eigenvalues - pca_number_of_modes: Number of PCA modes to use. Default: -1 (use all) + pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all) pca_template_model_point_subsample: Step size for subsampling model points. Default: 4 post_pca_transform: Optional ITK transform to apply after PCA registration. Default: None - fixed_distance_map: ITK image providing the distance map. + fixed_distance_map: ITK image providing the distance map, in mm. Default: None fixed_model: PyVista model used to compute the distance map, if one isn't provided. + Also supplies the target points for the symmetric term. reference_image: ITK image providing coordinate frame for computing the distance map. + pca_prior_weight: Weight (in mm) of the Mahalanobis shape prior + ``lambda * sum(b_i**2)``. Because b is expressed in standard + deviations, this term is the squared Mahalanobis distance in + shape space and makes the fit a MAP estimate rather than a pure + data fit constrained only by the coefficient bounds. + Default: 0.0 (prior disabled). + symmetric_weight: Weight in [0, 1] of the target-to-model distance + term. 0.0 measures model-to-target only, which lets the model + satisfy the metric while covering just part of the target. + Requires fixed_model; ignored with a warning when only a + distance map is available. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO Raises: - ValueError: If pca_eigenvector dimensions don't match model points + ValueError: If pca_eigenvector dimensions don't match model points, + if the mode counts disagree, or if neither a distance map nor a + fixed model plus reference image is provided. """ # Initialize base class with logging super().__init__(class_name="RegisterModelsPCA", log_level=log_level) # Store model data self.pca_template_model: pv.DataSet = pca_template_model - self.pca_eigenvectors: np.ndarray = pca_eigenvectors - self.pca_std_deviations: np.ndarray = pca_std_deviations + self.pca_eigenvectors: np.ndarray = np.asarray( + pca_eigenvectors, dtype=np.float64 + ) + self.pca_std_deviations: np.ndarray = np.asarray( + pca_std_deviations, dtype=np.float64 + ) + + if self.pca_eigenvectors.ndim != 2: + raise ValueError( + f"pca_eigenvectors must be 2D (modes, n_points*3), got shape " + f"{self.pca_eigenvectors.shape}" + ) + expected_size = pca_template_model.n_points * 3 + if self.pca_eigenvectors.shape[1] != expected_size: + raise ValueError( + f"Component dimension mismatch: expected {expected_size} " + f"(3 × {pca_template_model.n_points} points), got " + f"{self.pca_eigenvectors.shape[1]}" + ) + if self.pca_eigenvectors.shape[0] != self.pca_std_deviations.shape[0]: + raise ValueError( + f"Mode count mismatch: {self.pca_eigenvectors.shape[0]} eigenvectors " + f"but {self.pca_std_deviations.shape[0]} standard deviations" + ) self.post_pca_transform = post_pca_transform self._contour_tools = ContourTools() + self.fixed_model: Optional[pv.DataSet] = fixed_model self.fixed_distance_map = fixed_distance_map if ( self.fixed_distance_map is None and fixed_model is not None and reference_image is not None ): - self.fixed_model = fixed_model - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) elif self.fixed_distance_map is not None and ( fixed_model is not None or reference_image is not None ): self.log_warning( - "Fixed model and reference image will be ignored because a distance map is provided." + "A distance map was provided, so the reference image is ignored; " + "the fixed model is retained only for the symmetric metric term." ) elif self.fixed_distance_map is None and ( fixed_model is None or reference_image is None @@ -160,9 +213,11 @@ def __init__( self.pca_number_of_modes: int = pca_number_of_modes if self.pca_number_of_modes <= 0: - self.pca_number_of_modes = len(pca_std_deviations) + self.pca_number_of_modes = len(self.pca_std_deviations) self.pca_template_model_point_subsample = pca_template_model_point_subsample + self.pca_prior_weight = pca_prior_weight + self.symmetric_weight = symmetric_weight # outputs self.registered_model_pca_coefficients: Optional[np.ndarray] = None @@ -172,17 +227,26 @@ def __init__( self.forward_point_transform: Optional[itk.DisplacementFieldTransform] = None self.inverse_point_transform: Optional[itk.DisplacementFieldTransform] = None - # Image interpolator (created when needed) - self._fixed_distance_map_interpolator: Optional[ - itk.LinearInterpolateImageFunction - ] = None + # Sampling caches, built lazily by _prepare_sampling() + self._sampling_ready: bool = False + self._analytic_gradient: bool = True self._fixed_distance_map_max_distance: float = 0.0 + self._post_pca_affine_key: Optional[itk.Transform] = None + self._post_pca_affine: Optional[tuple[np.ndarray, np.ndarray]] = None self._metric_call_count: int = 0 - # Pre-convert mean shape points to ITK format - self._pca_template_model_points_itk: Optional[list[itk.Point]] = None - self._create_itk_points() + def _create_distance_map( + self, fixed_model: pv.DataSet, reference_image: itk.Image + ) -> itk.Image: + """Build the unsigned, un-normalized (mm) distance map of the target.""" + return self._contour_tools.create_distance_map( + fixed_model, + reference_image, + squared_distance=False, + negative_inside=False, + zero_inside=True, + ) @classmethod def from_json( @@ -195,6 +259,8 @@ def from_json( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from PCA model JSON file. @@ -217,6 +283,8 @@ def from_json( for registration. If None, must be set later before registration. fixed_model: Target surface mesh to register to. Default: None reference_image: Reference image defining coordinate space. Default: None + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO @@ -288,6 +356,8 @@ def from_json( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) @@ -302,6 +372,8 @@ def from_pca_model( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from a PCA model dictionary. @@ -320,6 +392,8 @@ def from_pca_model( fixed_distance_map: ITK image providing the distance values for registration. fixed_model: Target surface mesh to register to. reference_image: Reference image defining coordinate space. + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level. Returns: @@ -334,12 +408,6 @@ def from_pca_model( if "components" not in pca_model: raise ValueError("'components' field not found in pca_model") pca_eigenvectors = np.array(pca_model["components"], dtype=np.float64) - expected_size = pca_template_model.n_points * 3 - if pca_eigenvectors.shape[1] != expected_size: - raise ValueError( - f"Component dimension mismatch: expected {expected_size} " - f"(3 × {pca_template_model.n_points} points), got {pca_eigenvectors.shape[1]}" - ) return cls( pca_template_model=pca_template_model, pca_eigenvectors=pca_eigenvectors, @@ -350,38 +418,18 @@ def from_pca_model( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) - def _create_itk_points(self) -> None: - """Pre-convert mean shape points to ITK Point format for efficiency. - - This method creates ITK Point objects once at initialization, avoiding - repeated conversions during optimization iterations. - """ - self.log_info("Converting mean shape points to ITK format...") - - self._pca_template_model_points_itk = [] - for point in self.pca_template_model.points: - itk_point = itk.Point[itk.D, 3]() - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - self._pca_template_model_points_itk.append(itk_point) - - self.log_info( - f" Converted {len(self._pca_template_model_points_itk)} points to ITK format" - ) - def set_fixed_model( self, fixed_model: pv.UnstructuredGrid, reference_image: Optional[itk.Image] ) -> None: - """Set the fixed model for registration. - - If this is set, the fixed distance map will be set to None. + """Set the fixed model for registration and rebuild its distance map. Args: - fixed_model: PyVista model used to compute the distance map, if one isn't provided. + fixed_model: PyVista model used to compute the distance map. reference_image: ITK image providing coordinate frame for computing the distance map. """ if reference_image is None: @@ -389,26 +437,20 @@ def set_fixed_model( "reference_image must not be None when setting a fixed model" ) - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_model = fixed_model + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_fixed_distance_map(self, fixed_distance_map: Optional[itk.Image]) -> None: - """Set the reference image for registration. - - If this is set, the fixed model will be set to None. + """Set the distance map used as the registration target. Args: - fixed_distance_map: ITK image providing distance data + fixed_distance_map: ITK image providing distance data, in mm """ self.fixed_distance_map = fixed_distance_map - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> None: """Set the average model for registration. @@ -418,132 +460,298 @@ def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> Non (unstructured grid or polydata) """ self.pca_template_model = pca_template_model + self._sampling_ready = False + self.log_info(" Average model set successfully!") - self._pca_template_model_points_itk = None + def _affine_of_transform( + self, transform: itk.Transform + ) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Recover (matrix, offset) if ``transform`` acts affinely, else None. - self._create_itk_points() - self.log_info(" Average model set successfully!") + Probing the transform rather than querying ``GetMatrix()`` works for any + ITK transform type and correctly rejects the non-affine ones (such as a + displacement field), for which no constant Jacobian exists. + """ - def _mean_distance_metric( - self, - params: np.ndarray, - ) -> float: - """Evaluate the optimization metric (mean intensity) at model points. + def apply(vector: np.ndarray) -> np.ndarray: + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = (float(v) for v in vector) + result = transform.TransformPoint(point) + return np.array([result[0], result[1], result[2]], dtype=np.float64) - This is the objective function to be MAXIMIZED during optimization. - Higher values indicate better alignment with bright regions. + offset = apply(np.zeros(3)) + matrix = np.column_stack( + [apply(basis) - offset for basis in np.eye(3, dtype=np.float64)] + ) + probe = np.array([0.37, -0.61, 0.83], dtype=np.float64) + scale = max(1.0, float(np.abs(matrix).max()), float(np.abs(offset).max())) + if not np.allclose(apply(probe), matrix @ probe + offset, atol=1e-9 * scale): + return None + return matrix, offset - Args: - pca_deformation: Nx3 numpy array of PCA deformation vectors to add to points. - If None, no deformation is applied. + def _get_post_pca_affine(self) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Return the cached (matrix, offset) of post_pca_transform, or None. - Returns: - Mean distance value across all points + Keyed on the transform object so that reassigning post_pca_transform + invalidates the cache. """ - pca_deformation = self._compute_pca_deformation(params) - - # Create interpolator if not already cached (inline creation) - if self._fixed_distance_map_interpolator is None: - if self.fixed_distance_map is None: - self.log_error("Distance map is not set.") - raise ValueError("Distance map must be set before registering.") - ImageType = type(self.fixed_distance_map) - self._fixed_distance_map_interpolator = itk.LinearInterpolateImageFunction[ - ImageType, itk.D - ].New() - self._fixed_distance_map_interpolator.SetInputImage(self.fixed_distance_map) - fixed_distance_map_array = itk.GetArrayFromImage(self.fixed_distance_map) - self._fixed_distance_map_max_distance = fixed_distance_map_array.max() - self.log_debug("Interpolator created") - self.log_debug( - " Max distance = %s", self._fixed_distance_map_max_distance - ) + if self.post_pca_transform is None: + return None + if self._post_pca_affine_key is not self.post_pca_transform: + self._post_pca_affine_key = self.post_pca_transform + self._post_pca_affine = self._affine_of_transform(self.post_pca_transform) + return self._post_pca_affine + + def _prepare_sampling(self) -> None: + """Build the cached arrays the objective and its gradient are made of. + + Everything that does not depend on the PCA coefficients is computed once + here: the subsampled template points, the per-mode displacement vectors + already scaled by their standard deviation and mapped through the + post-PCA transform, the distance map and its gradient, and the + physical-to-index mapping used to sample them. + """ + if self.fixed_distance_map is None: + self.log_error("Distance map is not set.") + raise ValueError("Distance map must be set before registering.") + + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) + step = max(1, self.pca_template_model_point_subsample) + self._sample_slice = slice(None, None, step) + self._sample_points = template_points[self._sample_slice] + + # (modes, m, 3) displacement per unit coefficient, in template space. + modes = self.pca_eigenvectors.reshape(self.pca_eigenvectors.shape[0], -1, 3) + self._sample_modes = ( + modes[:, self._sample_slice, :] * self.pca_std_deviations[:, None, None] + ) - self.log_debug("Evaluating params = %s", params) - self.log_debug(" Max displacement = %s", pca_deformation.max(axis=0)) + # Fold the post-PCA transform into the mode directions so the gradient + # is expressed directly in world space. A non-affine post-PCA transform + # has no constant Jacobian, so the analytic gradient is disabled. + affine = self._get_post_pca_affine() + if self.post_pca_transform is not None and affine is None: + self.log_warning( + "post_pca_transform is not affine; falling back to a " + "finite-difference gradient." + ) + self._sample_modes_world = ( + self._sample_modes if affine is None else self._sample_modes @ affine[0].T + ) + self._analytic_gradient = self.post_pca_transform is None or affine is not None + + # Physical point -> continuous index: index = affine_inv @ (p - origin). + image = self.fixed_distance_map + direction = itk.array_from_matrix(image.GetDirection()) + index_to_world = direction @ np.diag(np.asarray(image.GetSpacing())) + self._index_to_world = index_to_world + self._world_to_index = np.linalg.inv(index_to_world) + self._image_origin = np.asarray(image.GetOrigin(), dtype=np.float64) + size = image.GetLargestPossibleRegion().GetSize() + self._image_size = np.array([size[0], size[1], size[2]], dtype=np.float64) + + # Array axes are (k, j, i), so index component a is array axis 2 - a. + # The forward differences are the exact derivative of the trilinear + # interpolant used to sample the map, which keeps the objective and its + # gradient consistent; a central difference would not. + self._distance_array = np.asarray( + itk.array_view_from_image(image), dtype=np.float64 + ) + self._fixed_distance_map_max_distance = float(self._distance_array.max()) + self._distance_forward_diff = tuple( + np.diff(self._distance_array, axis=2 - axis) + if self._distance_array.shape[2 - axis] > 1 + else None + for axis in range(3) + ) - # Sample distance at each point - n_valid_points = 0 - n_invalid_points = 0 - total_distance = 0.0 - center = np.zeros(3) - point = itk.Point[itk.D, 3]() - assert self.fixed_distance_map is not None, "fixed_distance_map must be set" - assert self._pca_template_model_points_itk is not None, ( - "ITK points must be initialized" + # Target points for the symmetric term. + self._target_points: Optional[np.ndarray] = None + if self.symmetric_weight > 0.0: + if self.fixed_model is None: + self.log_warning( + "symmetric_weight is %.3g but no fixed_model is available; " + "the target-to-model term is disabled.", + self.symmetric_weight, + ) + else: + self._target_points = np.asarray( + self.fixed_model.points, dtype=np.float64 + )[self._sample_slice] + + self._sampling_ready = True + self.log_debug( + "Sampling prepared: %d model points, %s target points, max distance %.3f mm", + self._sample_points.shape[0], + "no" if self._target_points is None else str(len(self._target_points)), + self._fixed_distance_map_max_distance, ) - image_size = self.fixed_distance_map.GetBufferedRegion().GetSize() - for i, base_point in enumerate(self._pca_template_model_points_itk): - if i % self.pca_template_model_point_subsample != 0: - continue - # Start with base point - point[0] = base_point[0] - point[1] = base_point[1] - point[2] = base_point[2] + def _sample_distance( + self, points: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, int]: + """Sample the distance map and its spatial gradient at world points. - # Add PCA deformation if provided - point[0] += pca_deformation[i, 0] - point[1] += pca_deformation[i, 1] - point[2] += pca_deformation[i, 2] + Points outside the image are clamped to the grid and charged the extra + travel from the clamped location, which keeps the metric continuous and + gives the optimizer a gradient that pushes such points back inside -- + unlike a constant out-of-bounds penalty, which is flat. - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) + Args: + points: (n, 3) array of world-space points - # Check if point is inside image bounds + Returns: + Tuple of (distances (n,), world-space gradients (n, 3), n_outside) + """ + index = (points - self._image_origin) @ self._world_to_index.T + clamped = np.clip(index, 0.0, self._image_size - 1.0) + outside = index - clamped + is_outside = np.any(outside != 0.0, axis=1) + n_outside = int(np.count_nonzero(is_outside)) + + # map_coordinates indexes the array as (k, j, i). + array_coordinates = clamped[:, ::-1] + distances = map_coordinates( + self._distance_array, array_coordinates.T, order=1, mode="nearest" + ) - coord_index = ( - self.fixed_distance_map.TransformPhysicalPointToContinuousIndex(point) + # d/dc_a of the trilinear interpolant is the forward difference along a, + # interpolated linearly across the other two axes within the same cell. + gradient_index = np.zeros_like(clamped) + for axis in range(3): + differences = self._distance_forward_diff[axis] + if differences is None: + continue + coordinates = array_coordinates.copy() + coordinates[:, 2 - axis] = np.clip( + np.floor(clamped[:, axis]), 0.0, self._image_size[axis] - 2.0 ) - if ( - 0 <= coord_index[0] < image_size[0] - and 0 <= coord_index[1] < image_size[1] - and 0 <= coord_index[2] < image_size[2] - ): - center[0] += point[0] - center[1] += point[1] - center[2] += point[2] - distance = ( - self._fixed_distance_map_interpolator.EvaluateAtContinuousIndex( - coord_index - ) - ) - total_distance += distance - n_valid_points += 1 - else: - n_invalid_points += 1 - - if n_invalid_points >= 0.05 * n_valid_points: - self.log_warning( - "%d of %d mapped outside of image. Rejecting.", - n_invalid_points, - n_valid_points + n_invalid_points, + gradient_index[:, axis] = map_coordinates( + differences, coordinates.T, order=1, mode="nearest" ) - return self._fixed_distance_map_max_distance + # A clamped axis cannot change the sampled value, so it carries no + # gradient from the map; the out-of-bounds term supplies it instead. + gradient_index[outside != 0.0] = 0.0 + gradients = gradient_index @ self._world_to_index + + if n_outside: + outside_world = outside @ self._index_to_world.T + outside_distance = np.linalg.norm(outside_world, axis=1) + safe = np.where(outside_distance > 0.0, outside_distance, 1.0) + distances = distances + outside_distance + gradients = gradients + outside_world / safe[:, None] + + return distances, gradients, n_outside + + def _deform(self, pca_coefficients: np.ndarray) -> np.ndarray: + """Deform the subsampled template points into world space.""" + n_modes = len(pca_coefficients) + points = self._sample_points + np.tensordot( + pca_coefficients, self._sample_modes[:n_modes], axes=(0, 0) + ) + return self._apply_post_pca_transform(points) - # Compute mean distance - mean_distance = total_distance / n_valid_points - center /= n_valid_points + def _objective_and_gradient(self, params: np.ndarray) -> tuple[float, np.ndarray]: + """Evaluate the registration objective and its gradient. - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - if log_level_int <= logging.DEBUG or self._metric_call_count % 100 == 0: - self.log_info( - " Metric %d: %s -> %f", - (self._metric_call_count + 1), - center, - mean_distance, + The objective is MINIMIZED: the distance map is zero on the target + surface and grows away from it. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Tuple of (objective value in mm, gradient with respect to params) + """ + if not self._sampling_ready: + self._prepare_sampling() + + n_modes = len(params) + modes_world = self._sample_modes_world[:n_modes] + points = self._deform(params) + + distances, gradients, n_outside = self._sample_distance(points) + forward_distance = float(distances.mean()) + # d/db_j of mean_i D(p_i) = mean_i grad_D(p_i) . (sigma_j * v_ij) + forward_gradient = np.einsum("ia,jia->j", gradients, modes_world) / len(points) + + weight = self.symmetric_weight if self._target_points is not None else 0.0 + reverse_distance = 0.0 + reverse_gradient = np.zeros(n_modes, dtype=np.float64) + if weight > 0.0: + assert self._target_points is not None, "target points must be set" + # Target -> model: each target point is charged its distance to the + # nearest deformed model point, so a model that covers only part of + # the target scores badly. The forward term alone cannot see this. + nearest, nearest_index = cKDTree(points).query(self._target_points) + nearest_distance = np.atleast_1d(np.asarray(nearest, dtype=np.float64)) + reverse_distance = float(nearest_distance.mean()) + safe = np.where(nearest_distance > 0.0, nearest_distance, 1.0) + direction = (points[nearest_index] - self._target_points) / safe[:, None] + # Accumulate each target's pull onto the model point it selected, + # then contract once against the modes. + pull = np.zeros_like(points) + np.add.at(pull, nearest_index, direction) + pull /= len(self._target_points) + reverse_gradient = np.einsum("ia,jia->j", pull, modes_world) + + objective = (1.0 - weight) * forward_distance + weight * reverse_distance + gradient = (1.0 - weight) * forward_gradient + weight * reverse_gradient + + prior = 0.0 + if self.pca_prior_weight > 0.0: + prior = self.pca_prior_weight * float(np.dot(params, params)) + objective += prior + gradient = gradient + 2.0 * self.pca_prior_weight * params + + if n_outside > 0.25 * len(points): + self.log_warning( + "%d of %d model points mapped outside the distance map.", + n_outside, + len(points), ) + + if self.log_level <= logging.DEBUG or self._metric_call_count % 25 == 0: self.log_info( - " Params %s", - params, + " Metric %d: %.4f mm (model->target %.4f, target->model %.4f, " + "prior %.4f, outside %d)", + self._metric_call_count + 1, + objective, + forward_distance, + reverse_distance, + prior, + n_outside, ) + self.log_debug(" Params %s", params) self._metric_call_count += 1 - return mean_distance + return objective, gradient + + def _mean_distance_metric(self, params: np.ndarray) -> float: + """Evaluate the registration objective at the given PCA coefficients. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Objective value, in mm. Lower is better. + """ + return self._objective_and_gradient(np.asarray(params, dtype=np.float64))[0] + + def _apply_post_pca_transform(self, points: np.ndarray) -> np.ndarray: + """Apply post_pca_transform to an (n, 3) array of world points.""" + affine = self._get_post_pca_affine() + if affine is not None: + return np.asarray(points @ affine[0].T + affine[1], dtype=np.float64) + if self.post_pca_transform is None: + return points + transformed = np.empty_like(points) + point = itk.Point[itk.D, 3]() + for i, source in enumerate(points): + point[0], point[1], point[2] = (float(v) for v in source) + result = self.post_pca_transform.TransformPoint(point) + transformed[i] = (result[0], result[1], result[2]) + return transformed def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: """Compute PCA deformation vectors for all points. @@ -552,28 +760,18 @@ def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: displacement = Σ(b_i * std_i * pca_eigenvector_i) Args: - pca_coefficients: Array of PCA coefficients b_i (one per mode) - pca_number_of_modes: Number of PCA modes to use. Default: use all available modes + pca_coefficients: Array of PCA coefficients b_i. Only as many modes + as there are coefficients contribute. Returns: - Nx3 array of deformation vectors (displacement from mean shape) + Nx3 array of deformation vectors (displacement from the template) """ - # Initialize deformation to zero - deformation = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - # Add contribution from each PCA mode - for i in range(self.pca_number_of_modes): - pca_eigenvector_flat = self.pca_eigenvectors[i, :] - - # Reshape to (N, 3) - pca_eigenvector_3d = pca_eigenvector_flat.reshape(-1, 3) - - # Add weighted deformation: b_i * std_i * pca_eigenvector_i - deformation += ( - pca_coefficients[i] * self.pca_std_deviations[i] * pca_eigenvector_3d - ) - - return deformation + n_modes = len(pca_coefficients) + scaled = pca_coefficients * self.pca_std_deviations[:n_modes] + deformation = np.asarray( + scaled @ self.pca_eigenvectors[:n_modes], dtype=np.float64 + ) + return deformation.reshape(-1, 3) def _optimize_pca_coefficients( self, @@ -584,12 +782,12 @@ def _optimize_pca_coefficients( ) -> tuple[np.ndarray, float]: """Optimize PCA coefficients - This method optimizes PCA mode coefficients to deform the model to better match - low values in the distance map. + Minimizes the mean distance between the deformed model and the target, + supplying the analytic gradient of the objective to the optimizer. Args: pca_number_of_modes: Number of PCA modes to use in optimization. Using fewer - modes provides smoother deformations. Default: 10 + modes provides smoother deformations. Default: 0 (use all) pca_coefficient_bounds: Bound on PCA coefficients in units of std deviations. Default: 3.0 (±3 std deviations per mode) method: Optimization method for scipy.optimize.minimize. @@ -600,52 +798,65 @@ def _optimize_pca_coefficients( Returns: Tuple of (pca_coefficients, mean_distance): - pca_coefficients: Optimized PCA coefficients - - mean_distance: Final mean distance metric value + - mean_distance: Final objective value, in mm Raises: ValueError: If number of PCA modes to use exceeds available modes """ + n_available = len(self.pca_eigenvectors) if pca_number_of_modes <= 0: - pca_number_of_modes = len(self.pca_eigenvectors) - if pca_number_of_modes > len(self.pca_eigenvectors): + pca_number_of_modes = n_available + if pca_number_of_modes > n_available: raise ValueError( - f"Number of PCA modes to use ({pca_number_of_modes}) exceeds available modes ({len(self.pca_std_deviations)})" + f"Number of PCA modes to use ({pca_number_of_modes}) exceeds " + f"available modes ({n_available})" ) self.pca_number_of_modes = pca_number_of_modes + self._prepare_sampling() + self.log_info(f"Number of PCA modes: {pca_number_of_modes}") self.log_info( f"PCA coefficient bounds: ±{pca_coefficient_bounds} std deviations" ) self.log_info(f"Optimization method: {method}") self.log_info(f"Max iterations: {max_iterations}") + self.log_info(f"Shape prior weight: {self.pca_prior_weight}") + self.log_info(f"Symmetric weight: {self.symmetric_weight}") - bounds = [] - for _ in range(pca_number_of_modes): - bounds.append((-pca_coefficient_bounds, pca_coefficient_bounds)) + bounds = [ + (-pca_coefficient_bounds, pca_coefficient_bounds) + for _ in range(pca_number_of_modes) + ] - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - disp = log_level_int <= logging.INFO + disp = self.log_level <= logging.INFO + + # The metric is in mm, so the default gradient tolerance is meaningful. + # Without an analytic gradient the finite-difference step must be large + # enough to move sample points by a useful fraction of a voxel. + options: dict = {"maxiter": max_iterations, "disp": disp, "gtol": 1e-6} + if not self._analytic_gradient: + options["eps"] = 1e-2 self.log_info("Running optimization...") result_pca = minimize( # type: ignore[call-overload] - lambda params: self._mean_distance_metric(params), - np.zeros(self.pca_number_of_modes), + self._objective_and_gradient + if self._analytic_gradient + else self._mean_distance_metric, + np.zeros(pca_number_of_modes), method=method, + jac=self._analytic_gradient, bounds=bounds, - options={"maxiter": max_iterations, "disp": disp}, + options=options, ) optimized_pca_coefficients = result_pca.x - optimized_mean_distance = result_pca.fun + optimized_mean_distance = float(result_pca.fun) self.log_info("Optimization completed!") self.log_info(f"Optimized PCA coefficients: {optimized_pca_coefficients}") - self.log_info(f"Final mean intensity: {optimized_mean_distance:.2f}") + self.log_info(f"Metric evaluations: {self._metric_call_count}") + self.log_info(f"Final mean distance: {optimized_mean_distance:.4f} mm") return optimized_pca_coefficients, optimized_mean_distance @@ -672,35 +883,12 @@ def transform_template_model(self) -> pv.DataSet: self.registered_model_pca_coefficients, ) - # Apply deformation and affine transform to each point - final_points = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - n_points = self.pca_template_model.n_points - progress_interval = max(1, n_points // 10) # Report progress every 10% - - point = itk.Point[itk.D, 3]() - for i in range(n_points): - # Report progress - if i % progress_interval == 0 or i == n_points - 1: - self.log_progress(i + 1, n_points, prefix="Transforming points") - - # Start with mean shape point - point[0] = float(self.pca_template_model.points[i][0]) - point[1] = float(self.pca_template_model.points[i][1]) - point[2] = float(self.pca_template_model.points[i][2]) - - # Add PCA deformation - point[0] += self.registered_model_pca_deformation[i, 0] - point[1] += self.registered_model_pca_deformation[i, 1] - point[2] += self.registered_model_pca_deformation[i, 2] - - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) - - # Store result - final_points[i, 0] = point[0] - final_points[i, 1] = point[1] - final_points[i, 2] = point[2] + # Deform in the template frame, then map into target space. + deformed_points = ( + np.asarray(self.pca_template_model.points, dtype=np.float64) + + self.registered_model_pca_deformation + ) + final_points = self._apply_post_pca_transform(deformed_points) # Create new model with transformed points self.registered_model = self.pca_template_model.copy(deep=True) @@ -717,30 +905,36 @@ def transform_point( point: itk.Point, include_post_pca_transform: bool = True, ) -> itk.Point: - """Transform an arbitrary point using nearest neighbor interpolation. + """Transform an arbitrary point through the PCA deformation field. Args: point: ITK point to transform (itk.Point[itk.D, 3]) + include_post_pca_transform: Also apply post_pca_transform. Default: True Returns: Transformed ITK point + Raises: + ValueError: If compute_pca_transforms() has not been called yet + Notes: - 1) if the point is outside the image bounds, the point is not transformed. - 2) if the forward point transform is set, it is applied. - 3) if the post_pca_transform is set and enabled, it is applied. - 4) if the forward point transform is not set, no errors are raised. + This samples the *approximated* deformation field built by + compute_pca_transforms(), which is splatted and blurred, so it does + not reproduce transform_template_model() exactly; the RMS of that + difference is logged when the field is built. Points outside the + field's reference image are not displaced. Example: >>> p = itk.Point[itk.D, 3]() >>> p[0], p[1], p[2] = 10.0, 20.0, 30.0 >>> transformed_p = registrar.transform_point(p) """ - - if self.forward_point_transform is not None: - transformed_point = self.forward_point_transform.TransformPoint(point) - else: - transformed_point = point + if self.forward_point_transform is None: + self.log_error("Forward point transform is not set.") + raise ValueError( + "compute_pca_transforms() must be called before transform_point()" + ) + transformed_point = self.forward_point_transform.TransformPoint(point) if include_post_pca_transform and self.post_pca_transform is not None: transformed_point = self.post_pca_transform.TransformPoint( @@ -749,9 +943,21 @@ def transform_point( return transformed_point - def compute_pca_transforms(self, reference_image: itk.Image) -> dict: + def compute_pca_transforms( + self, reference_image: itk.Image, blur_sigma: float = 2.5 + ) -> dict: """Compute PCA transforms. + The field is built by splatting the per-point PCA displacements onto the + reference grid and blurring them, so it only approximates the exact + per-point deformation. The RMS of that approximation error, and of the + forward/inverse round trip, are both logged. + + Args: + reference_image: ITK image providing the coordinate frame for the field. + blur_sigma: Sigma for Gaussian blurring of the deformation field. + Default: 2.5 + Returns: Dictionary containing: - 'forward_point_transform': POINT transform mapping template @@ -761,17 +967,19 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: Note: These are point transforms, oriented opposite to image-registration - transforms; see docs/developer/transform_conventions. + transforms; see docs/developer/transform_conventions. Neither + includes post_pca_transform. """ assert self.registered_model_pca_deformation is not None, ( "PCA deformation must be computed" ) + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) template_model_pca_deformation_field_image = ( self._contour_tools.create_deformation_field( - np.array(self.pca_template_model.points), + template_points, self.registered_model_pca_deformation, reference_image=reference_image, - blur_sigma=2.5, + blur_sigma=blur_sigma, ptype=itk.D, ) ) @@ -787,11 +995,44 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: self.forward_point_transform ) ) + + self._log_transform_fidelity(template_points) + return { "forward_point_transform": self.forward_point_transform, "inverse_point_transform": self.inverse_point_transform, } + def _log_transform_fidelity(self, template_points: np.ndarray) -> None: + """Report how well the field reproduces the deformation and inverts.""" + assert self.forward_point_transform is not None, "forward transform must be set" + assert self.inverse_point_transform is not None, "inverse transform must be set" + assert self.registered_model_pca_deformation is not None, ( + "PCA deformation must be computed" + ) + + point = itk.Point[itk.D, 3]() + forward = np.empty_like(template_points) + round_trip = np.empty_like(template_points) + for i, source in enumerate(template_points): + point[0], point[1], point[2] = (float(v) for v in source) + mapped = self.forward_point_transform.TransformPoint(point) + forward[i] = (mapped[0], mapped[1], mapped[2]) + back = self.inverse_point_transform.TransformPoint(mapped) + round_trip[i] = (back[0], back[1], back[2]) + + expected = template_points + self.registered_model_pca_deformation + field_rms = float(np.sqrt(np.mean(np.sum((forward - expected) ** 2, axis=1)))) + inverse_rms = float( + np.sqrt(np.mean(np.sum((round_trip - template_points) ** 2, axis=1))) + ) + self.log_info( + "Deformation field RMS error: %.4f mm (approximation of the " + "per-point deformation)", + field_rms, + ) + self.log_info("Forward/inverse round-trip RMS error: %.4f mm", inverse_rms) + def register( self, pca_number_of_modes: int = 0, @@ -799,40 +1040,40 @@ def register( method: str = "L-BFGS-B", max_iterations: int = 100, ) -> dict: - """Optimize PCA coefficients to deform the model to better match - low values in the distance map. + """Optimize PCA coefficients to deform the model onto the target. Args: pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all available modes) - pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.0 + pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.5 method: Optimization method for scipy.optimize.minimize. Default: 'L-BFGS-B' (supports bounds) max_iterations: Maximum number of optimization iterations. - Default: 50 + Default: 100 Returns: Dictionary containing: - 'registered_model': Final registered PyVista model - 'pca_coefficients': Optimized PCA coefficients - - 'mean_distance': Final mean distance metric value + - 'mean_distance': Final objective value, in mm Raises: - ValueError: If reference image is not set + ValueError: If the distance map is not set Example: >>> result = registrar.register(pca_number_of_modes=10) >>> result['registered_model'].save('registered_heart.vtk') """ if self.fixed_distance_map is None: - raise ValueError("Reference image must be set before registration") + raise ValueError("A distance map must be set before registration") if pca_number_of_modes <= 0: pca_number_of_modes = self.pca_number_of_modes - self.log_section("PCA-BASED MODEL-TO-IMAGE REGISTRATION", width=70) + self.log_section("PCA-BASED MODEL-TO-MODEL REGISTRATION", width=70) self.log_info(f"Number of points: {self.pca_template_model.n_points}") self.log_info(f"Modes to use: {pca_number_of_modes}") + self._metric_call_count = 0 self.registered_model_pca_coefficients, self.registered_model_mean_distance = ( self._optimize_pca_coefficients( pca_number_of_modes=pca_number_of_modes, diff --git a/tests/test_register_models_pca.py b/tests/test_register_models_pca.py index 66b19746..2ce2fbcd 100644 --- a/tests/test_register_models_pca.py +++ b/tests/test_register_models_pca.py @@ -8,11 +8,12 @@ import numpy as np import pytest import pyvista as pv +from scipy.optimize import approx_fprime from physiotwin4d.register_models_pca import RegisterModelsPCA -def _make_registrar() -> RegisterModelsPCA: +def _make_registrar(**kwargs: Any) -> RegisterModelsPCA: """Create a small PCA registrar with a three-point template surface.""" template_model = pv.PolyData( np.array( @@ -27,27 +28,52 @@ def _make_registrar() -> RegisterModelsPCA: pca_eigenvectors = np.zeros((1, template_model.n_points * 3), dtype=np.float64) pca_std_deviations = np.ones(1, dtype=np.float64) fixed_distance_map = itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)) + kwargs.setdefault("symmetric_weight", 0.0) return RegisterModelsPCA( pca_template_model=template_model, pca_eigenvectors=pca_eigenvectors, pca_std_deviations=pca_std_deviations, pca_number_of_modes=1, fixed_distance_map=fixed_distance_map, + **kwargs, ) -def test_itk_template_points_are_distinct_objects() -> None: - """Cached ITK points are distinct per template vertex.""" - registrar = _make_registrar() +def _sphere_registrar( + radius: float, + modes: int = 1, + **kwargs: Any, +) -> tuple[RegisterModelsPCA, np.ndarray]: + """Build a registrar whose single mode inflates a sphere radially. + + Returns the registrar and the per-point unit-norm eigenvector it was given. + """ + template = pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24) + directions = np.asarray(template.points, dtype=np.float64) + directions /= np.linalg.norm(directions, axis=1, keepdims=True) + + eigenvectors = np.zeros((modes, template.n_points * 3), dtype=np.float64) + eigenvectors[0] = directions.reshape(-1) / np.linalg.norm(directions) + for mode in range(1, modes): + # Orthogonal filler modes: displace a disjoint slab of points along x. + filler = np.zeros((template.n_points, 3), dtype=np.float64) + filler[mode::modes, 0] = 1.0 + eigenvectors[mode] = filler.reshape(-1) / np.linalg.norm(filler) - points = registrar._pca_template_model_points_itk - assert points is not None - assert len({id(point) for point in points}) == len(points) - assert [float(points[0][0]), float(points[1][0]), float(points[2][1])] == [ - 0.0, - 1.0, - 1.0, - ] + reference_image = itk.image_from_array(np.zeros((48, 48, 48), dtype=np.float32)) + reference_image.SetSpacing([1.0, 1.0, 1.0]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=template, + pca_eigenvectors=eigenvectors, + pca_std_deviations=np.full(modes, 5.0), + pca_template_model_point_subsample=1, + fixed_model=pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24), + reference_image=reference_image, + **kwargs, + ) + return registrar, eigenvectors[0].reshape(-1, 3) def test_set_fixed_model_requires_reference_image() -> None: @@ -60,6 +86,48 @@ def test_set_fixed_model_requires_reference_image() -> None: ) +def test_mode_count_mismatch_is_rejected() -> None: + """Eigenvector and standard-deviation counts must agree.""" + template_model = pv.PolyData(np.zeros((3, 3), dtype=np.float64)) + with pytest.raises(ValueError, match="Mode count mismatch"): + RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=np.zeros((2, 9), dtype=np.float64), + pca_std_deviations=np.ones(3, dtype=np.float64), + fixed_distance_map=itk.image_from_array( + np.zeros((4, 4, 4), dtype=np.float32) + ), + ) + + +def test_compute_pca_deformation_scales_eigenvectors_by_std() -> None: + """Deformation is exactly sum(b_i * std_i * eigenvector_i), reshaped (N, 3).""" + template_model = pv.PolyData(np.zeros((2, 3), dtype=np.float64)) + eigenvectors = np.array( + [ + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float64, + ) + std_deviations = np.array([2.0, 5.0], dtype=np.float64) + registrar = RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=eigenvectors, + pca_std_deviations=std_deviations, + fixed_distance_map=itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)), + ) + + deformation = registrar._compute_pca_deformation(np.array([1.5, -1.0])) + + # Point 0: 1.5*2*[1,0,0] + (-1)*5*[0,0,1]; point 1: 1.5*2*[0,1,0]. + assert np.allclose(deformation, [[3.0, 0.0, -5.0], [0.0, 3.0, 0.0]]) + + # A shorter coefficient vector uses only the leading modes. + leading = registrar._compute_pca_deformation(np.array([1.5])) + assert np.allclose(leading, [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0]]) + + def test_transform_template_model_applies_post_pca_transform_after_deformation() -> ( None ): @@ -78,3 +146,184 @@ def test_transform_template_model_applies_post_pca_transform_after_deformation() assert np.allclose(registered_model.points[0], [2.0, 0.0, 0.0]) assert np.allclose(registered_model.points[1], [4.0, 0.0, 0.0]) + + +def test_modes_are_deformed_in_the_template_frame_then_transformed() -> None: + """Regression: modes must be rotated with the template, not added after it. + + The registered model must equal ``A @ (template + deformation)``, never + ``A @ template + deformation``. The two differ whenever the post-PCA + transform contains a rotation, which is the case for every ICP alignment. + """ + registrar = _make_registrar() + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + deformation = np.tile( + np.array([1.0, 0.0, 0.0], dtype=np.float64), + (registrar.pca_template_model.n_points, 1), + ) + registrar.registered_model_pca_deformation = deformation + + # 90 degrees about z, so x-displacements must come out along y. + matrix = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64 + ) + offset = np.array([10.0, -3.0, 2.0], dtype=np.float64) + transform = itk.AffineTransform[itk.D, 3].New() + transform.SetMatrix(itk.matrix_from_array(matrix)) + transform.SetTranslation(offset) + registrar.post_pca_transform = transform + + registered_model: Any = registrar.transform_template_model() + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + correct = (template_points + deformation) @ matrix.T + offset + wrong = template_points @ matrix.T + offset + deformation + + assert np.allclose(registered_model.points, correct) + assert not np.allclose(registered_model.points, wrong) + + +def test_analytic_gradient_matches_finite_differences() -> None: + """The supplied Jacobian agrees with a finite-difference gradient.""" + registrar, _ = _sphere_registrar(radius=8.0, modes=3, symmetric_weight=0.5) + registrar.pca_prior_weight = 0.1 + registrar._prepare_sampling() + + params = np.array([0.4, -0.25, 0.15], dtype=np.float64) + _, analytic = registrar._objective_and_gradient(params) + numeric = approx_fprime(params, registrar._mean_distance_metric, 1e-5) + + assert np.allclose(analytic, numeric, atol=2e-3) + + +def test_register_recovers_known_coefficients() -> None: + """Fitting to a target built from known coefficients recovers them.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + # Target = template inflated by b = 1.0 along the single radial mode. + truth = 1.0 + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + truth * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + result = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert result["pca_coefficients"][0] == pytest.approx(truth, abs=0.1) + assert result["mean_distance"] < 0.5 + + +def test_symmetric_term_penalizes_partial_coverage() -> None: + """The target-to-model term sees coverage the model-to-target term misses. + + A hemisphere sitting on a full sphere scores near-perfectly one-way: every + model point lies on the target surface. Only the target-to-model term + notices that half the target has no model near it. + """ + target = pv.Sphere(radius=8.0, theta_resolution=24, phi_resolution=24) + points = np.asarray(target.points, dtype=np.float64) + hemisphere = pv.PolyData(points[points[:, 2] > 0.0]) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=hemisphere, + pca_eigenvectors=np.zeros((1, hemisphere.n_points * 3), dtype=np.float64), + pca_std_deviations=np.ones(1, dtype=np.float64), + pca_template_model_point_subsample=1, + fixed_model=target, + reference_image=reference_image, + symmetric_weight=0.0, + ) + + registrar._prepare_sampling() + one_way = registrar._mean_distance_metric(np.zeros(1)) + + registrar.symmetric_weight = 0.5 + registrar._prepare_sampling() + symmetric = registrar._mean_distance_metric(np.zeros(1)) + + # Model points all lie on the target surface, so the one-way term is small. + assert one_way < 0.5 + assert symmetric > 1.0 + + +def test_prior_shrinks_coefficients() -> None: + """Raising pca_prior_weight pulls the solution toward the mean shape.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + unregularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + registrar.pca_prior_weight = 5.0 + registrar._sampling_ready = False + regularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert abs(regularized["pca_coefficients"][0]) < abs( + unregularized["pca_coefficients"][0] + ) + + +def test_transform_point_requires_computed_transforms() -> None: + """transform_point raises instead of silently returning the input.""" + registrar = _make_registrar() + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = 1.0, 2.0, 3.0 + + with pytest.raises(ValueError, match="compute_pca_transforms"): + registrar.transform_point(point) + + +def test_pca_transforms_round_trip() -> None: + """forward reproduces the deformation and inverse undoes it.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + registrar.registered_model_pca_deformation = ( + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + transforms = registrar.compute_pca_transforms(reference_image, blur_sigma=1.5) + forward = transforms["forward_point_transform"] + inverse = transforms["inverse_point_transform"] + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + expected = template_points + registrar.registered_model_pca_deformation + + point = itk.Point[itk.D, 3]() + mapped = np.empty_like(template_points) + back = np.empty_like(template_points) + for i, source in enumerate(template_points): + point[0], point[1], point[2] = (float(v) for v in source) + forward_point = forward.TransformPoint(point) + mapped[i] = (forward_point[0], forward_point[1], forward_point[2]) + inverse_point = inverse.TransformPoint(forward_point) + back[i] = (inverse_point[0], inverse_point[1], inverse_point[2]) + + # The field is splatted and blurred, so it only approximates the deformation. + field_rms = np.sqrt(np.mean(np.sum((mapped - expected) ** 2, axis=1))) + round_trip_rms = np.sqrt(np.mean(np.sum((back - template_points) ** 2, axis=1))) + + assert field_rms < 1.5 + assert round_trip_rms < 1.0 From f78e9a813cd53222271ea79e04511e0ae7eeaf62 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 06:39:31 -0400 Subject: [PATCH 2/5] BUG: Convert Greedy transforms from RAS to LPS, unify registration seeding Greedy reports its affine in RAS while ITK is LPS, but the 4x4 was copied straight into an itk.AffineTransform, negating x and y. Recovering a known (6, -4, 3) mm shift returned (+5.96, -4.01, -2.92) instead of (-6, +4, -3), and warping by the result scored below the unregistered pair (foreground NCC 0.21 vs 0.32). Rigid, Affine and Deformable were all affected; the displacement field was already LPS and is left alone. Existing Greedy tests only asserted the transforms were non-None, so the sign error survived. - Change basis in RegisterImagesGreedy._matrix_to_itk_affine - Add known-shift accuracy tests for Greedy, ANTs and ICON (a KnownShiftCase helper in conftest); ANTs and ICON were audited and are correct - Replace per-backend initial_forward_transform handling, which pre-warped in ANTs, pre-warped only the image in ICON, and double-applied in Greedy, with one RegisterImagesBase.register_from() - Add TransformTools.invert_transform, preferring the analytic inverse over a displacement field that is only defined on the reference grid - Remove prior_weight from RegisterTimeSeriesImages, the reconstruction workflow and its CLI flag - Offer explicit Rigid/Similarity/Affine modes in RegisterModelsICP, with bounding-box scale estimation before ICP - Add ImageTools.pad_image; drop icon_iterations from RegisterModelsDistanceMaps.register() - Fix PhaseSampleDataset caching (0 now means unbounded, as documented) and include Case8Deploy in tutorial 09's case discovery - Regenerate registration_time_series_images baselines --- docs/developer/registration_images.rst | 25 +- .../Heart-Create_Statistical_Model/README.md | 3 +- .../cli/reconstruct_highres_4d_ct.py | 13 - src/physiotwin4d/image_tools.py | 144 +++++++++- src/physiotwin4d/physicsnemo_tools.py | 2 +- src/physiotwin4d/register_images_ants.py | 71 +---- src/physiotwin4d/register_images_base.py | 155 +++++++++- src/physiotwin4d/register_images_chain.py | 52 ++-- src/physiotwin4d/register_images_greedy.py | 139 ++++----- src/physiotwin4d/register_images_icon.py | 29 -- .../register_models_distance_maps.py | 84 ++++-- src/physiotwin4d/register_models_icp.py | 267 +++++++++++++----- .../register_time_series_images.py | 113 +------- src/physiotwin4d/train_physicsnemo_mgn.py | 16 ++ src/physiotwin4d/transform_tools.py | 62 +++- ...rkflow_fit_statistical_model_to_patient.py | 83 ++++-- .../workflow_reconstruct_highres_4d_ct.py | 21 -- .../basic_forward_transform_0.hdf | 4 +- .../basic_time_series_registered_0.mha | 4 +- .../middle_frame_forward_transform_0.hdf | 3 + .../prior_forward_transform_0.hdf | 3 - .../prior_time_series_registered_0.mha | 3 - .../transform_application_time_series_0.mha | 4 +- tests/conftest.py | 73 +++++ tests/test_register_images_ants.py | 53 +++- tests/test_register_images_chain.py | 46 ++- tests/test_register_images_greedy.py | 44 +++ tests/test_register_images_icon.py | 37 ++- tests/test_register_time_series_images.py | 50 +--- ...torial_06_lung_create_statistical_model.py | 8 +- ...7_lung_fit_statistical_model_to_patient.py | 14 +- .../tutorial_09_lung_train_physicsnemo_mgn.py | 48 +++- 32 files changed, 1099 insertions(+), 574 deletions(-) create mode 100644 tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf delete mode 100644 tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf delete mode 100644 tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha diff --git a/docs/developer/registration_images.rst b/docs/developer/registration_images.rst index dd649dd3..e3fdf8d8 100644 --- a/docs/developer/registration_images.rst +++ b/docs/developer/registration_images.rst @@ -57,10 +57,10 @@ Workflows that accept a ``registration_method`` (e.g. :class:`WorkflowConvertImageToUSD`, :class:`RegisterTimeSeriesImages`) take any :class:`RegisterImagesBase` instance, including a composite chain that runs multiple backends in sequence. :class:`RegisterImagesChain` runs an -ordered list of registrars, feeding each stage's ``forward_transform`` as the -next stage's ``initial_forward_transform``. :class:`RegisterImagesGreedyICON` -is a named 2-stage convenience class for the common case of a fast Greedy -registration followed by ICON refinement: +ordered list of registrars, each stage refining the previous stage's +``forward_transform`` through ``register_from()`` (see `Seeding a registration`_ +below). :class:`RegisterImagesGreedyICON` is a named 2-stage convenience class +for the common case of a fast Greedy registration followed by ICON refinement: .. code-block:: python @@ -76,6 +76,23 @@ registration followed by ICON refinement: registrar.greedy.set_number_of_iterations([30, 15, 7, 3]) registrar.icon.set_number_of_iterations(20) +Seeding a registration +====================== + +To start from an alignment you already have, call ``register_from()`` instead of +``register()``: + +.. code-block:: python + + result = registrar.register_from(known_forward_transform, moving_image) + +It warps the moving image, mask and labelmap onto the fixed grid by that +transform, registers the residual, and composes the two, so the returned +transforms still map between the *original* moving image and the fixed image. +Every backend goes through this one implementation -- no registrar accepts an +initial transform of its own, which is what keeps the pre-warp, the composition +and the inversion identical for Greedy, ICON and ANTs. + Development Notes ================= diff --git a/experiments/Heart-Create_Statistical_Model/README.md b/experiments/Heart-Create_Statistical_Model/README.md index 92727e9d..d450fe00 100644 --- a/experiments/Heart-Create_Statistical_Model/README.md +++ b/experiments/Heart-Create_Statistical_Model/README.md @@ -228,7 +228,8 @@ For ICON registration: - Check alignment quality from step 2 (ICP should produce good initial alignment) - Verify average surface looks reasonable before step 3 - If Greedy affine fails, check input mesh quality and topology -- If ICON deformable quality is poor, increase `icon_iterations` in the `register()` call +- If ICON deformable quality is poor, increase the iteration count on the + registrar's ICON instance: `registrar.registrar_ICON.set_number_of_iterations(...)` ### Import Errors - Ensure all PhysioTwin4D dependencies are installed diff --git a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py index fbb13d8c..c44fc805 100644 --- a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py @@ -103,13 +103,6 @@ def main() -> int: "(default: enabled; use --no-register-reference for an identity transform)" ), ) - parser.add_argument( - "--prior-weight", - type=float, - default=0.0, - help="Weight for temporal smoothing with prior transforms [0.0-1.0] (default: 0.0)", - ) - # Registration iterations parser.add_argument( "--Greedy-iterations", @@ -221,11 +214,6 @@ def main() -> int: ) return 1 - # Validate prior weight - if not 0.0 <= args.prior_weight <= 1.0: - print(f"Error: Prior weight must be in [0.0, 1.0], got {args.prior_weight}") - return 1 - # Create output directory os.makedirs(args.output_dir, exist_ok=True) @@ -296,7 +284,6 @@ def main() -> int: # Configure registration parameters workflow.set_modality(args.modality) - workflow.set_prior_weight(args.prior_weight) workflow.set_mask_dilation(args.mask_dilation_mm) if fixed_mask is not None: diff --git a/src/physiotwin4d/image_tools.py b/src/physiotwin4d/image_tools.py index 0e5c4e22..284c71fb 100644 --- a/src/physiotwin4d/image_tools.py +++ b/src/physiotwin4d/image_tools.py @@ -6,7 +6,7 @@ """ import logging -from typing import Any, Optional, Union, overload +from typing import Any, Optional, Union, cast, overload import itk import numpy as np @@ -270,6 +270,148 @@ def make_isotropic_image(self, image: itk.Image) -> itk.Image: result.DisconnectPipeline() return result + @staticmethod + def _per_axis_values( + value: Union[float, int, list, tuple, NDArray[Any]], + dimension: int, + name: str, + ) -> list[float]: + """Broadcast a scalar to every axis, or validate a per-axis sequence. + + Args: + value: Scalar applied to every axis, or one value per axis. + dimension: Number of image dimensions expected. + name: Parameter name, used in error messages. + + Returns: + One value per axis. + + Raises: + ValueError: If a sequence has the wrong length, or any value is + negative. + """ + if np.isscalar(value): + values = [float(cast(float, value))] * dimension + else: + values = [float(v) for v in value] # type: ignore[union-attr] + if len(values) != dimension: + raise ValueError( + f"{name} needs a scalar or one value per image dimension " + f"({dimension}), got {len(values)}." + ) + if any(v < 0.0 for v in values): + raise ValueError(f"{name} must be >= 0, got {values}") + return values + + def pad_image( + self, + image: itk.Image, + pad_portion: Optional[ + Union[float, list[float], tuple[float, ...], NDArray[Any]] + ] = None, + pad_voxels: Optional[ + Union[int, list[int], tuple[int, ...], NDArray[Any]] + ] = None, + background_value: float = 0.0, + ) -> itk.Image: + """Pad *image* on every side with a constant-valued margin. + + The margin is given either as a portion of each axis' physical extent + (*pad_portion*) or directly in voxels (*pad_voxels*); exactly one of the + two must be supplied. Either accepts a scalar, applied to every axis, or + one value per image dimension. Both pad the lower and the upper end of + every axis, so ``pad_voxels=10`` grows all six faces of a 3-D image by + ten voxels. Spacing and direction are untouched. + + The origin and size are updated together so the original voxels keep + their physical positions: the padded image's index ``(0, 0, 0)`` sits one + margin below the input's, and the input data occupies the interior. + (``itk.ConstantPadImageFilter`` alone reports the margin as a negative + start index instead, which most file formats drop on write — shifting the + data. The region-of-interest pass here folds that index back into the + origin.) + + Args: + image: ITK image to pad. + pad_portion: Portion of an axis' physical extent (``size * spacing``) + to add at both ends, as a scalar for every axis or one value per + dimension: ``0.1`` grows an axis spanning 200 mm by 20 mm per + side. Rounded up to whole voxels. Mutually exclusive with + *pad_voxels*. + pad_voxels: Margin in voxels, as a scalar for every axis or one value + per dimension, applied at both ends of each axis. Mutually + exclusive with *pad_portion*. + background_value: Pixel value written into the new margin + (default: 0.0). + + Returns: + Padded image with the same pixel type, spacing and direction. + + Raises: + ValueError: If neither or both of *pad_portion* and *pad_voxels* are + given, if either is negative, or if a sequence does not have one + entry per image dimension. + """ + if (pad_portion is None) == (pad_voxels is None): + raise ValueError( + "Specify exactly one of pad_portion or pad_voxels; got " + f"pad_portion={pad_portion}, pad_voxels={pad_voxels}." + ) + + size = [int(s) for s in image.GetLargestPossibleRegion().GetSize()] + if pad_portion is not None: + portions = self._per_axis_values(pad_portion, len(size), "pad_portion") + # extent_i = size_i * spacing_i, and pad_portion * extent_i of margin + # is that distance divided by spacing_i, so the spacing cancels. + margin = [int(np.ceil(p * s)) for p, s in zip(portions, size)] + self.log_info( + "Padding by %s voxels per side (%s of extent); size %s -> %s", + margin, + [f"{p * 100.0:.1f}%" for p in portions], + size, + [s + 2 * p for s, p in zip(size, margin)], + ) + else: + assert pad_voxels is not None # guaranteed by the check above + margin = [ + int(v) + for v in self._per_axis_values(pad_voxels, len(size), "pad_voxels") + ] + self.log_info( + "Padding by %s voxels per side; size %s -> %s", + margin, + size, + [s + 2 * p for s, p in zip(size, margin)], + ) + + ImageType = type(image) + # SetConstant is typed to the pixel type, so an integer image rejects a + # Python float. + pixel_type = itk.template(image)[1][0] + constant = ( + float(background_value) + if pixel_type in (itk.F, itk.D) + else int(round(background_value)) + ) + + pad_filter = itk.ConstantPadImageFilter[ImageType, ImageType].New() + pad_filter.SetInput(image) + pad_filter.SetPadLowerBound(margin) + pad_filter.SetPadUpperBound(margin) + pad_filter.SetConstant(constant) + pad_filter.Update() + padded = pad_filter.GetOutput() + + # Re-anchor the padded region at index 0, moving the margin into the + # origin so the physical position of the original data is preserved. + roi_filter = itk.RegionOfInterestImageFilter[ImageType, ImageType].New() + roi_filter.SetInput(padded) + roi_filter.SetRegionOfInterest(padded.GetLargestPossibleRegion()) + roi_filter.Update() + result = roi_filter.GetOutput() + result.DisconnectPipeline() + return result + def binary_dilate_image( self, image: itk.Image, diff --git a/src/physiotwin4d/physicsnemo_tools.py b/src/physiotwin4d/physicsnemo_tools.py index 19a1c3f1..346d9865 100644 --- a/src/physiotwin4d/physicsnemo_tools.py +++ b/src/physiotwin4d/physicsnemo_tools.py @@ -405,8 +405,8 @@ def _target_values(self, path: Path) -> np.ndarray: f"{path} has {values.shape[0]} points, expected {self._n_points}." ) + self._cache[path] = values if self._cache_max_samples != 0: - self._cache[path] = values while len(self._cache) > self._cache_max_samples: self._cache.popitem(last=False) return values diff --git a/src/physiotwin4d/register_images_ants.py b/src/physiotwin4d/register_images_ants.py index 5c2f013e..a19df20b 100644 --- a/src/physiotwin4d/register_images_ants.py +++ b/src/physiotwin4d/register_images_ants.py @@ -510,7 +510,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using ANTs registration algorithm. @@ -524,14 +523,6 @@ def registration_method( region of interest in the moving image moving_image_pre (itk.Image, optional): Pre-processed moving image. If None, preprocessing is performed automatically - initial_forward_transform (itk.Transform, optional): Initial - forward transform (same convention as the returned - forward_transform: used to warp the moving image onto the fixed - grid). Can be any ITK transform type (Affine, Rigid, - DisplacementField, Composite, etc.). It is applied by pre-warping - the moving image onto the fixed grid before registration; the - returned transforms compose this initial alignment with the - registration refinement. Returns: dict: Dictionary containing: @@ -573,9 +564,9 @@ def registration_method( >>> registrar.set_fixed_mask(heart_mask_fixed) >>> result = registrar.register(moving_image, moving_mask=heart_mask_moving) >>> - >>> # Registration with initial transform + >>> # Registration seeded with a known alignment >>> initial_tfm = itk.AffineTransform[itk.D, 3].New() - >>> result = registrar.register(moving_image, initial_forward_transform=initial_tfm) + >>> result = registrar.register_from(initial_tfm, moving_image) """ if moving_image is not None: self.moving_image = moving_image @@ -591,22 +582,6 @@ def registration_method( if self.fixed_image_pre is None: self.fixed_image_pre = self.preprocess(self.fixed_image, self.modality) - if initial_forward_transform is not None: - self.log_info("Pre-warping moving image with initial transform...") - transform_tools = TransformTools() - self.moving_image_pre = transform_tools.transform_image( - self.moving_image_pre, - initial_forward_transform, - self.fixed_image, - ) - if self.moving_mask is not None: - self.moving_mask = transform_tools.transform_image( - self.moving_mask, - initial_forward_transform, - self.fixed_image, - interpolation_method="nearest", - ) - transform_type = None if self.transform_type == "Deformable": transform_type = "antsRegistrationSyNQuick[so]" @@ -706,46 +681,8 @@ def registration_method( reference_image=self.moving_image, ) - # Important: ANTs does NOT include the initial_transform in the output transforms - # We need to manually compose them - if initial_forward_transform is not None: - self.log_info("Composing initial transform with registration result...") - - # For forward_transform (Moving -> Fixed): Apply initial_forward_transform first, then registration - # Transform order: point -> initial_forward_transform -> forward_reg - forward_transform = itk.CompositeTransform[itk.D, 3].New() - forward_transform.AddTransform(initial_forward_transform) - # Add transforms from registration result (may be composite) - if isinstance(forward_reg, itk.CompositeTransform[itk.D, 3]): - for i in range(forward_reg.GetNumberOfTransforms()): - forward_transform.AddTransform(forward_reg.GetNthTransform(i)) - else: - forward_transform.AddTransform(forward_reg) - - # For inverse_transform (Fixed -> Moving): Apply registration inverse first, then initial inverse - # Transform order: point -> inverse_reg -> initial_forward_transform^(-1) - inverse_transform = itk.CompositeTransform[itk.D, 3].New() - # Add registration inverse transforms - if isinstance(inverse_reg, itk.CompositeTransform[itk.D, 3]): - for i in range(inverse_reg.GetNumberOfTransforms()): - inverse_transform.AddTransform(inverse_reg.GetNthTransform(i)) - else: - inverse_transform.AddTransform(inverse_reg) - # Invert and add initial transform - # For displacement field transforms, we need to invert properly - transform_tools = TransformTools() - initial_inverse = transform_tools.invert_displacement_field_transform( - transform_tools.convert_transform_to_displacement_field_transform( - initial_forward_transform, self.moving_image - ) - ) - inverse_transform.AddTransform(initial_inverse) - - self.log_info("Transforms composed successfully") - else: - # No initial transform, use registration results directly - forward_transform = forward_reg - inverse_transform = inverse_reg + forward_transform = forward_reg + inverse_transform = inverse_reg moving_image_reg = registration_result["warpedmovout"] loss = ants.image_similarity( self._itk_to_ants_image(self.fixed_image), diff --git a/src/physiotwin4d/register_images_base.py b/src/physiotwin4d/register_images_base.py index 64e0d7c1..ab496adf 100644 --- a/src/physiotwin4d/register_images_base.py +++ b/src/physiotwin4d/register_images_base.py @@ -236,7 +236,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Main registration method to align moving image to fixed image. @@ -252,7 +251,6 @@ def registration_method( moving_mask (itk.image, optional): Binary mask for moving image ROI moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Preprocessed moving image - initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed Returns: dict: Dictionary containing: @@ -274,7 +272,6 @@ def register( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register a moving image to the fixed image. @@ -282,12 +279,14 @@ def register( concrete subclasses. It should align the moving image to the fixed image using the specific algorithm implemented by the subclass. + To start from a known alignment, use :meth:`register_from` rather than + seeding the backend directly. + Args: moving_image (itk.image): The 3D image to be registered to the fixed image moving_mask (itk.image, optional): Binary mask for moving image ROI moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Preprocessed moving image - initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed Returns: dict: Dictionary containing transformation results: @@ -347,7 +346,6 @@ def register( moving_mask=new_moving_mask, moving_labelmap=moving_labelmap, moving_image_pre=moving_image_pre, - initial_forward_transform=initial_forward_transform, ) self.forward_transform = result["forward_transform"] @@ -360,6 +358,153 @@ def register( "loss": self.loss, } + def register_from( + self, + initial_forward_transform: itk.Transform, + moving_image: itk.Image, + moving_mask: Optional[itk.Image] = None, + moving_labelmap: Optional[itk.Image] = None, + ) -> dict[str, Union[itk.Transform, float]]: + """Register starting from a known alignment. + + The moving data is warped onto the fixed grid by + ``initial_forward_transform`` first, :meth:`register` then measures only + the residual misalignment, and the two are composed. This is the single + supported way to seed a registration: doing it here rather than inside + each backend keeps the pre-warp, the composition and the inversion + identical for every algorithm. + + The image, the mask and the labelmap are all pre-warped, so they stay in + the same frame as each other; the mask and labelmap use nearest-neighbor + interpolation to preserve their discrete values. + + Args: + initial_forward_transform: Starting alignment, in the same + convention as the returned ``forward_transform`` -- it warps the + moving image onto the fixed grid. + moving_image: The 3D image to be registered to the fixed image. + moving_mask: Binary mask for the moving image ROI. + moving_labelmap: Multi-label segmentation for the moving image. + + Returns: + dict: Same keys as :meth:`register`, with the transforms composed so + they map between the *original* moving image and the fixed image. + + Raises: + ValueError: If the fixed image has not been set. + """ + warped_image, warped_mask, warped_labelmap = self._prewarp_moving( + initial_forward_transform, moving_image, moving_mask, moving_labelmap + ) + result = self.register( + warped_image, + moving_mask=warped_mask, + moving_labelmap=warped_labelmap, + ) + composed = self._compose_with_initial( + initial_forward_transform, result, moving_image + ) + + self.forward_transform = composed["forward_transform"] + self.inverse_transform = composed["inverse_transform"] + self.loss = composed["loss"] + return composed + + def _prewarp_moving( + self, + initial_forward_transform: itk.Transform, + moving_image: itk.Image, + moving_mask: Optional[itk.Image], + moving_labelmap: Optional[itk.Image], + ) -> tuple[itk.Image, Optional[itk.Image], Optional[itk.Image]]: + """Warp the moving image, mask and labelmap onto the fixed grid. + + Args: + initial_forward_transform: Alignment to apply, in the image-warp + convention. + moving_image: Raw moving image. + moving_mask: Moving mask, or None. + moving_labelmap: Moving labelmap, or None. + + Returns: + Tuple of the warped ``(image, mask, labelmap)``, the latter two None + when not supplied. The mask and labelmap are warped with + nearest-neighbor interpolation to keep their discrete values. + + Raises: + ValueError: If the fixed image has not been set. + """ + if self.fixed_image is None: + raise ValueError("Fixed image must be set before registration.") + + transform_tools = TransformTools() + self.log_info("Pre-warping moving data with the initial transform...") + + def _warp(image: Optional[itk.Image], nearest: bool) -> Optional[itk.Image]: + if image is None: + return None + return transform_tools.transform_image( + image, + initial_forward_transform, + self.fixed_image, + interpolation_method="nearest" if nearest else "linear", + ) + + return ( + _warp(moving_image, nearest=False), + _warp(moving_mask, nearest=True), + _warp(moving_labelmap, nearest=True), + ) + + def _compose_with_initial( + self, + initial_forward_transform: itk.Transform, + result: dict[str, Union[itk.Transform, float]], + moving_image: itk.Image, + ) -> dict[str, Union[itk.Transform, float]]: + """Compose a residual registration result onto its initial transform. + + Args: + initial_forward_transform: The alignment the moving data was + pre-warped by. + result: Result of registering the pre-warped data. + moving_image: Raw moving image, whose grid defines the domain the + initial transform is inverted over. + + Returns: + The result dict with both transforms mapping between the *original* + moving image and the fixed image. + """ + transform_tools = TransformTools() + + # The registration measured the residual from the pre-warped position, + # so the total is the initial transform followed by that residual. An + # itk.CompositeTransform applies its transforms in reverse order of + # addition, so adding the initial first makes the residual apply first -- + # which is what the image-warp direction needs: a fixed-grid sample is + # mapped by the residual, then by the initial transform, to land in the + # original moving image. + forward_transform = itk.CompositeTransform[itk.D, 3].New() + forward_transform.AddTransform(initial_forward_transform) + forward_transform.AddTransform(cast(itk.Transform, result["forward_transform"])) + + # The inverse runs the other way -- a moving-grid sample is mapped by the + # initial transform's inverse into the pre-warped frame, then by the + # residual's inverse into the fixed image -- so the additions are + # reversed too. + initial_inverse = transform_tools.invert_transform( + initial_forward_transform, moving_image + ) + inverse_transform = itk.CompositeTransform[itk.D, 3].New() + inverse_transform.AddTransform(cast(itk.Transform, result["inverse_transform"])) + inverse_transform.AddTransform(initial_inverse) + + return { + "forward_transform": forward_transform, + "inverse_transform": inverse_transform, + "loss": result["loss"], + } + def _delegate_to( self, other: "RegisterImagesBase", diff --git a/src/physiotwin4d/register_images_chain.py b/src/physiotwin4d/register_images_chain.py index 0e5fab50..d05a719a 100644 --- a/src/physiotwin4d/register_images_chain.py +++ b/src/physiotwin4d/register_images_chain.py @@ -14,8 +14,9 @@ class RegisterImagesChain(RegisterImagesBase): - """Run an ordered list of registrars in sequence, feeding each stage's - forward_transform as the next stage's initial_forward_transform. + """Run an ordered list of registrars in sequence, each stage refining the + previous stage's forward_transform via + :meth:`RegisterImagesBase.register_from`. Use this to combine independent registration backends into a multi-stage pipeline (e.g. a fast coarse registrar followed by a refinement stage). @@ -67,12 +68,14 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Run each registrar in ``self.registrars`` in order. - Each stage's ``forward_transform`` becomes the next stage's - ``initial_forward_transform``. + The first stage registers the raw moving image; every later stage sees + the moving data pre-warped by the running result and contributes only a + refinement, which is composed back on -- the same mechanics as + :meth:`RegisterImagesBase.register_from`, run through the delegated + ``registration_method`` path so masks are not re-converted per stage. Note: ``moving_image_pre`` is ignored: each stage may need different @@ -87,25 +90,40 @@ def registration_method( moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Ignored - see Note above - initial_forward_transform (itk.Transform, optional): Initial - transformation from moving to fixed, used to initialize the - first stage Returns: dict: The last stage's result dict (see :meth:`RegisterImagesBase.register`) """ - current_initial = initial_forward_transform + current_forward: Optional[itk.Transform] = None result: dict[str, Union[itk.Transform, float]] = {} for registrar in self.registrars: - self._delegate_to(registrar, moving_image, moving_mask, moving_labelmap) - result = registrar.registration_method( - moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, + if current_forward is None: + stage_image, stage_mask, stage_labelmap = ( + moving_image, + moving_mask, + moving_labelmap, + ) + else: + stage_image, stage_mask, stage_labelmap = self._prewarp_moving( + current_forward, moving_image, moving_mask, moving_labelmap + ) + + self._delegate_to(registrar, stage_image, stage_mask, stage_labelmap) + stage_result = registrar.registration_method( + stage_image, + moving_mask=stage_mask, + moving_labelmap=stage_labelmap, moving_image_pre=None, - initial_forward_transform=current_initial, ) - self._capture_delegate_result(registrar, result) - current_initial = cast(itk.Transform, result["forward_transform"]) + self._capture_delegate_result(registrar, stage_result) + + result = ( + stage_result + if current_forward is None + else self._compose_with_initial( + current_forward, stage_result, moving_image + ) + ) + current_forward = cast(itk.Transform, result["forward_transform"]) return result diff --git a/src/physiotwin4d/register_images_greedy.py b/src/physiotwin4d/register_images_greedy.py index d5af2108..cd710c44 100644 --- a/src/physiotwin4d/register_images_greedy.py +++ b/src/physiotwin4d/register_images_greedy.py @@ -237,13 +237,38 @@ def _write_affine_matrix_file(self, mat_4x4: NDArray[np.float64]) -> str: self.log_debug("Wrote Greedy affine init matrix to %s", path) return path + # RAS <-> LPS basis change: negate x and y, leave z. Its own inverse. + _RAS_TO_LPS = np.diag([-1.0, -1.0, 1.0]) + def _matrix_to_itk_affine(self, mat_4x4: NDArray[np.float64]) -> itk.Transform: - """Convert 4x4 affine matrix to ITK AffineTransform.""" + """Convert Greedy's 4x4 RAS affine matrix to an ITK AffineTransform. + + Greedy works in RAS (its logs print "Final RAS Transform" and its + ``.mat`` files are RAS), while ITK -- and every transform this project + stores or applies -- is LPS. The matrix therefore has to change basis on + the way across: ``M_lps = F M_ras F`` and ``t_lps = F t_ras`` with + ``F = diag(-1, -1, 1)``. Skipping this negates the x and y components of + every Greedy result, which for a near-identity registration is a + sub-millimetre error that is easy to miss and impossible to correct + downstream. + + Only the affine needs this. Greedy's displacement fields come back as + images carrying the input's own LPS metadata, so + :meth:`_sitk_warp_to_itk_displacement_transform` passes them through + unchanged. See docs/developer/transform_conventions. + + Args: + mat_4x4: 4x4 affine matrix in Greedy's RAS convention. + + Returns: + The equivalent ITK affine transform, in LPS. + """ mat_4x4 = np.asarray(mat_4x4, dtype=np.float64) if mat_4x4.shape != (4, 4): raise ValueError(f"Expected 4x4 matrix, got shape {mat_4x4.shape}") - M = mat_4x4[:3, :3] - t = mat_4x4[:3, 3] + flip = self._RAS_TO_LPS + M = flip @ mat_4x4[:3, :3] @ flip + t = flip @ mat_4x4[:3, 3] center = itk.Point[itk.D, 3]() for i in range(3): center[i] = 0.0 @@ -283,7 +308,6 @@ def _registration_method_affine_or_rigid( moving_mask_sitk: Optional[Any] = None, fixed_labelmap_sitk: Optional[Any] = None, moving_labelmap_sitk: Optional[Any] = None, - initial_affine: Optional[NDArray[np.float64]] = None, ) -> tuple[NDArray[np.float64], float]: """Run Greedy affine or rigid registration. Returns (4x4 matrix, loss).""" Greedy3D = _try_import_greedy() @@ -307,19 +331,8 @@ def _registration_method_affine_or_rigid( cmd += " -gm fixed_mask -mm moving_mask" kwargs["fixed_mask"] = fixed_mask_sitk kwargs["moving_mask"] = moving_mask_sitk - # Greedy crashes (heap corruption) when an initial affine is passed as an - # in-memory matrix; write it to a temp file and pass the path instead. - initial_affine_file: Optional[str] = None - if initial_affine is not None: - initial_affine_file = self._write_affine_matrix_file(initial_affine) - cmd += f" -ia {initial_affine_file}" - self.log_debug("Greedy affine/rigid command: %s", cmd) - try: - g.execute(cmd, **kwargs) - finally: - if initial_affine_file is not None: - os.remove(initial_affine_file) + g.execute(cmd, **kwargs) mat = np.array(g["aff_out"], dtype=np.float64) try: ml = g.metric_log() @@ -339,36 +352,35 @@ def _registration_method_deformable( moving_mask_sitk: Optional[Any] = None, fixed_labelmap_sitk: Optional[Any] = None, moving_labelmap_sitk: Optional[Any] = None, - initial_affine: Optional[NDArray[np.float64]] = None, ) -> tuple[Optional[NDArray[np.float64]], Any, float]: """Run Greedy deformable registration. Returns (affine 4x4 or None, warp_sitk, loss).""" Greedy3D = _try_import_greedy() g = Greedy3D() - # Optional affine init (uses configured metric) - if initial_affine is None: - cmd_aff = "-d 3" - if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: - cmd_aff += " -w 0.60" - cmd_aff += " -i fixed moving" - kwargs_aff = { - "fixed": fixed_sitk, - "moving": moving_sitk, - } - if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: - cmd_aff += " -w 0.40 -i fixed_labelmap moving_labelmap" - kwargs_aff["fixed_labelmap"] = fixed_labelmap_sitk - kwargs_aff["moving_labelmap"] = moving_labelmap_sitk - cmd_aff += f" -a -dof 12 -n {iterations_str} -m {metric_str} -o aff_init" - kwargs_aff["aff_init"] = None - if fixed_mask_sitk is not None and moving_mask_sitk is not None: - cmd_aff += " -gm fixed_mask -mm moving_mask" - kwargs_aff["fixed_mask"] = fixed_mask_sitk - kwargs_aff["moving_mask"] = moving_mask_sitk - self.log_debug("Greedy deformable affine-init command: %s", cmd_aff) - g.execute(cmd_aff, **kwargs_aff) - initial_affine = np.array(g["aff_init"], dtype=np.float64) - self.log_info("Greedy deformable affine init complete") + # Greedy seeds its own deformable stage with an affine pass, using the + # configured metric. + cmd_aff = "-d 3" + if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: + cmd_aff += " -w 0.60" + cmd_aff += " -i fixed moving" + kwargs_aff = { + "fixed": fixed_sitk, + "moving": moving_sitk, + } + if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: + cmd_aff += " -w 0.40 -i fixed_labelmap moving_labelmap" + kwargs_aff["fixed_labelmap"] = fixed_labelmap_sitk + kwargs_aff["moving_labelmap"] = moving_labelmap_sitk + cmd_aff += f" -a -dof 12 -n {iterations_str} -m {metric_str} -o aff_init" + kwargs_aff["aff_init"] = None + if fixed_mask_sitk is not None and moving_mask_sitk is not None: + cmd_aff += " -gm fixed_mask -mm moving_mask" + kwargs_aff["fixed_mask"] = fixed_mask_sitk + kwargs_aff["moving_mask"] = moving_mask_sitk + self.log_debug("Greedy deformable affine-init command: %s", cmd_aff) + g.execute(cmd_aff, **kwargs_aff) + initial_affine = np.array(g["aff_init"], dtype=np.float64) + self.log_info("Greedy deformable affine init complete") # Greedy crashes (heap corruption) when the affine init is passed as an # in-memory matrix via -it; write it to a temp file and pass the path. @@ -415,13 +427,11 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using Greedy. Converts ITK images to SimpleITK, runs Greedy (affine and/or deformable), - then converts outputs back to ITK transforms. Composes with - initial_forward_transform when provided. + then converts outputs back to ITK transforms. Returns a dict with "forward_transform", "inverse_transform", and "loss". As with the other image-registration backends, @@ -487,24 +497,6 @@ def registration_method( iterations_str = self._greedy_iterations_str() metric_str = self._greedy_metric() - # Optional initial transform: convert ITK -> 4x4 for Greedy - initial_affine: Optional[NDArray[np.float64]] = None - if initial_forward_transform is not None: - # If it's affine-like, extract 4x4; else convert to displacement and skip for Greedy init - if hasattr(initial_forward_transform, "GetMatrix"): - M = np.eye(4, dtype=np.float64) - M[:3, :3] = np.asarray(initial_forward_transform.GetMatrix()).reshape( - 3, 3 - ) - if hasattr(initial_forward_transform, "GetTranslation"): - M[:3, 3] = np.asarray(initial_forward_transform.GetTranslation()) - if hasattr(initial_forward_transform, "GetCenter"): - c = np.asarray(initial_forward_transform.GetCenter()) - M[:3, 3] += c - M[:3, :3] @ c - initial_affine = M - # Non-affine initial: we could convert to disp field and pass; for simplicity we skip Greedy init - # and compose at the end (same as ANTs). - forward_transform: itk.Transform inverse_transform: itk.Transform loss_val: float @@ -520,7 +512,6 @@ def registration_method( iterations_str=iterations_str, metric_str=metric_str, dof=6, - initial_affine=initial_affine, ) forward_transform = self._matrix_to_itk_affine(mat) inverse_affine = itk.AffineTransform[itk.D, 3].New() @@ -537,7 +528,6 @@ def registration_method( iterations_str=iterations_str, metric_str=metric_str, dof=12, - initial_affine=initial_affine, ) forward_transform = self._matrix_to_itk_affine(mat) inverse_affine = itk.AffineTransform[itk.D, 3].New() @@ -554,7 +544,6 @@ def registration_method( moving_labelmap_sitk=moving_labelmap_sitk, iterations_str=iterations_str, metric_str=metric_str, - initial_affine=initial_affine, ) aff_tfm = ( self._matrix_to_itk_affine(aff_mat) if aff_mat is not None else None @@ -597,26 +586,6 @@ def registration_method( inverse_composite.AddTransform(inv_aff) inverse_transform = inverse_composite - # Compose with user-provided initial transform (same semantics as ANTs) - if initial_forward_transform is not None: - transform_tools = TransformTools() - forward_composite = itk.CompositeTransform[itk.D, 3].New() - forward_composite.AddTransform(initial_forward_transform) - forward_composite.AddTransform(forward_transform) - initial_disp = ( - transform_tools.convert_transform_to_displacement_field_transform( - initial_forward_transform, self.moving_image - ) - ) - inv_initial = transform_tools.invert_displacement_field_transform( - initial_disp - ) - inverse_composite = itk.CompositeTransform[itk.D, 3].New() - inverse_composite.AddTransform(inverse_transform) - inverse_composite.AddTransform(inv_initial) - forward_transform = forward_composite - inverse_transform = inverse_composite - return { "forward_transform": forward_transform, "inverse_transform": inverse_transform, diff --git a/src/physiotwin4d/register_images_icon.py b/src/physiotwin4d/register_images_icon.py index 1c964081..6b71a627 100644 --- a/src/physiotwin4d/register_images_icon.py +++ b/src/physiotwin4d/register_images_icon.py @@ -18,7 +18,6 @@ import numpy as np from .register_images_base import RegisterImagesBase -from .transform_tools import TransformTools DEFAULT_FINETUNE_LEARNING_RATE = 2e-5 @@ -196,7 +195,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using ICON registration algorithm. @@ -212,9 +210,6 @@ def registration_method( fixed_mask, enables mask-constrained registration moving_image_pre (itk.image, optional): Pre-processed moving image. If None, preprocessing is performed automatically - initial_forward_transform (itk.Transform, optional): Initial transformation from moving - to fixed. If provided, it is used to transform the moving image before - registration. Returns: dict: Dictionary containing: @@ -252,18 +247,10 @@ def registration_method( >>> result = registrar.register(moving_image, moving_mask=heart_mask_moving) """ - tfm_tools = TransformTools() - if moving_image_pre is None: moving_image_pre = self.preprocess(moving_image, self.modality) new_moving_image_pre = moving_image_pre - if initial_forward_transform is not None: - new_moving_image_pre = tfm_tools.transform_image( - moving_image_pre, - initial_forward_transform, - self.fixed_image, - ) # Prefer labelmap over binary mask when both sides have a labelmap. use_labelmaps = moving_labelmap is not None and self.fixed_labelmap is not None @@ -301,22 +288,6 @@ def registration_method( loss = loss_artifacts[0] - if initial_forward_transform is not None: - forward_transform = tfm_tools.combine_displacement_field_transforms( - initial_forward_transform, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=1.0, - mode="compose", - ) - - dftfm = tfm_tools.convert_transform_to_displacement_field_transform( - forward_transform, - self.fixed_image, - ) - inverse_transform = tfm_tools.invert_displacement_field_transform(dftfm) - return { "forward_transform": forward_transform, "inverse_transform": inverse_transform, diff --git a/src/physiotwin4d/register_models_distance_maps.py b/src/physiotwin4d/register_models_distance_maps.py index f7d94d83..4e863949 100644 --- a/src/physiotwin4d/register_models_distance_maps.py +++ b/src/physiotwin4d/register_models_distance_maps.py @@ -35,7 +35,7 @@ ... reference_image=reference_image, ... mask_dilation_mm=20, ... ) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=50) + >>> result = registrar.register(transform_type='Deformable') >>> >>> # Access results >>> aligned_model = result['registered_model'] @@ -86,6 +86,7 @@ class RegisterModelsDistanceMaps(PhysioTwin4DBase): fixed_model (pv.PolyData): Target surface model reference_image (itk.Image): Reference image for coordinate frame mask_dilation_mm (float): Dilation amount in mm for binary registration masks + distance_squared_max (float): Maximum squared distance for distance map normalization transform_tools (TransformTools): Transform utility instance contour_tools (ContourTools): Model utility instance registrar_Greedy (RegisterImagesGreedy): Greedy registration instance @@ -110,7 +111,7 @@ class RegisterModelsDistanceMaps(PhysioTwin4DBase): >>> result = registrar.register(transform_type='Affine') >>> >>> # Or run deformable (Greedy affine + ICON) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=50) + >>> result = registrar.register(transform_type='Deformable') >>> >>> # Get aligned model and transforms >>> aligned_model = result['registered_model'] @@ -122,6 +123,7 @@ def __init__( moving_model: pv.PolyData, fixed_model: pv.PolyData, reference_image: itk.Image, + distance_squared_max: float = 50.0, mask_dilation_mm: float = 20, log_level: int | str = logging.INFO, ): @@ -145,6 +147,7 @@ def __init__( self.moving_model = moving_model self.fixed_model = fixed_model self.reference_image = reference_image + self.distance_squared_max = distance_squared_max self.mask_dilation_mm = mask_dilation_mm # Utilities @@ -189,19 +192,31 @@ def _create_masks_from_models(self) -> None: squared_distance=True, negative_inside=True, zero_inside=False, - norm_to_max_distance=50.0, + norm_to_max_distance=self.distance_squared_max, ) - # 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 + if self.mask_dilation_mm > 0: + # 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_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm + ) + else: + self.fixed_mask_image = None + + itk.imwrite( + self.fixed_mask_image, "debug_fixed_mask_image.nii.gz", compression=True ) - self.fixed_mask_image = self.labelmap_tools.convert_labelmap_to_mask( - binary_mask, dilation_in_mm=self.mask_dilation_mm + itk.imwrite( + self.fixed_distance_map_image, + "debug_fixed_distance_map_image.nii.gz", + compression=True, ) # Create moving distance map @@ -211,7 +226,7 @@ def _create_masks_from_models(self) -> None: squared_distance=True, negative_inside=True, zero_inside=False, - norm_to_max_distance=50.0, + norm_to_max_distance=self.distance_squared_max, ) # Emulate CT intensity range by multiplying by 1000 @@ -221,16 +236,28 @@ def _create_masks_from_models(self) -> None: tmp_arr = itk.GetArrayViewFromImage(self.moving_distance_map_image) tmp_arr *= 1000 - # 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 + if self.mask_dilation_mm > 0: + # 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_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm + ) + else: + self.moving_mask_image = None + + itk.imwrite( + self.moving_mask_image, "debug_moving_mask_image.nii.gz", compression=True ) - self.moving_mask_image = self.labelmap_tools.convert_labelmap_to_mask( - binary_mask, dilation_in_mm=self.mask_dilation_mm + itk.imwrite( + self.moving_distance_map_image, + "debug_moving_distance_map_image.nii.gz", + compression=True, ) self.log_info("Distance map and mask generation complete") @@ -238,7 +265,6 @@ def _create_masks_from_models(self) -> None: def register( self, transform_type: str = "Deformable", - icon_iterations: int = 50, ) -> dict: """Perform mask-based registration of moving model to fixed model. @@ -259,7 +285,6 @@ def register( Args: transform_type: Registration transform type - 'None', 'Rigid', 'Affine', or 'Deformable'. Default: 'Deformable' - icon_iterations: Number of ICON optimization iterations for 'Deformable' mode. Default: 50 Returns: Dictionary containing: @@ -278,7 +303,7 @@ def register( >>> result = registrar.register(transform_type='Affine') >>> >>> # Deformable registration (Greedy affine + ICON) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=100) + >>> result = registrar.register(transform_type='Deformable') """ if transform_type not in ["None", "Rigid", "Affine", "Deformable"]: raise ValueError( @@ -319,10 +344,7 @@ def register( # Step 3: ICON deformable stage (only for Deformable mode) if transform_type == "Deformable": - self.log_info( - "Performing ICON deformable registration (%d iterations)...", - icon_iterations, - ) + self.log_info("Performing ICON deformable registration...") # Pre-align moving distance map and binary mask into the fixed grid using the Greedy affine result moving_distance_map_affine_transformed = ( @@ -340,8 +362,8 @@ def register( # interpolation_method="nearest", # ) - # Configure and run ICON - self.registrar_ICON.set_number_of_iterations(icon_iterations) + # Configure and run ICON. Iteration count and any other ICON tuning + # come from registrar_ICON itself, configured by the caller. self.registrar_ICON.set_fixed_image(self.fixed_distance_map_image) # self.registrar_ICON.set_fixed_mask(self.fixed_mask_image) diff --git a/src/physiotwin4d/register_models_icp.py b/src/physiotwin4d/register_models_icp.py index c845a658..e0602086 100644 --- a/src/physiotwin4d/register_models_icp.py +++ b/src/physiotwin4d/register_models_icp.py @@ -3,15 +3,17 @@ This module provides the RegisterModelsICP class for aligning anatomical models using Iterative Closest Point (ICP) algorithm. The workflow includes: 1. Initial centroid alignment -2. Rigid or affine ICP alignment +2. Isotropic bounding-box scaling (Similarity and Affine transform types) +3. Rigid, similarity or affine ICP alignment The registration is particularly useful for initial rough alignment of generic models to patient-specific anatomical data. Key Features: - Centroid-based initial alignment - - VTK ICP algorithm with rigid or affine transformation modes - - Three-stage affine pipeline: centroid → rigid ICP → affine ICP + - Bounding-box size matching before any ICP iteration + - VTK ICP with rigid, similarity or affine transformation modes + - Four-stage affine pipeline: centroid → scale → similarity ICP → affine ICP - Support for PyVista models - Automatic transform composition @@ -57,8 +59,15 @@ class RegisterModelsICP(PhysioTwin4DBase): centroid alignment for initialization followed by VTK's ICP algorithm. **Registration Pipelines:** - - **Rigid transform type**: Centroid alignment → Rigid ICP - - **Affine transform type**: Centroid alignment → Rigid ICP → Affine ICP + - **Rigid**: Centroid alignment → Rigid ICP + - **Similarity**: Centroid alignment → bounding-box scaling → + Similarity ICP + - **Affine**: Centroid alignment → bounding-box scaling → + Similarity ICP → Affine ICP + + The bounding-box scaling estimates the size ratio between the models from + their bounding-box diagonals before ICP starts. It is skipped for 'Rigid', + whose transform has no scale degree of freedom. **Transform Convention:** These are POINT transforms (applied with TransformPoint, e.g. via @@ -133,6 +142,118 @@ def __init__( self.inverse_point_transform: Optional[itk.AffineTransform] = None self.registered_model: Optional[pv.PolyData] = None + # ICP stages run for each transform type, in order. + _ICP_STAGES = { + "Rigid": ("Rigid",), + "Similarity": ("Similarity",), + "Affine": ("Similarity", "Affine"), + } + + def _icp_stage( + self, model: pv.PolyData, mode: str, max_iterations: int + ) -> tuple[pv.PolyData, itk.AffineTransform]: + """Run one VTK ICP stage against the fixed model. + + Args: + model: Current state of the moving model. + mode: Landmark-transform mode, one of ``"Rigid"``, ``"Similarity"`` + or ``"Affine"``. + max_iterations: Maximum ICP iterations for this stage. + + Returns: + Tuple of the transformed model and this stage's moving→fixed point + transform. + """ + icp = vtk.vtkIterativeClosestPointTransform() + icp.SetSource(model) + icp.SetTarget(self.fixed_model) + landmark = icp.GetLandmarkTransform() + if mode == "Rigid": + landmark.SetModeToRigidBody() + elif mode == "Similarity": + landmark.SetModeToSimilarity() + else: + landmark.SetModeToAffine() + icp.SetMaximumNumberOfIterations(max_iterations) + icp.Update() + + stage_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( + icp.GetMatrix() + ) + transformed = self.transform_tools.transform_pvcontour( + model, + stage_transform, + with_deformation_magnitude=False, + ) + return transformed, stage_transform + + def _bounding_box_scale( + self, moving_model: pv.PolyData, fixed_model: pv.PolyData + ) -> float: + """Return the isotropic scale matching the models' bounding-box diagonals. + + The diagonal is used rather than the three side lengths separately: the + bounding boxes are axis-aligned, so per-axis ratios would fold any + residual rotation between the models into an anisotropic stretch. A + single factor changes size without distorting shape, leaving the + remaining anisotropy to the affine ICP stage. + + Args: + moving_model: Model whose size is being matched. + fixed_model: Model supplying the target size. + + Returns: + The scale factor, or ``1.0`` when either model is degenerate (a + single point or a zero-extent bounding box), where no meaningful + ratio exists. + """ + + def _diagonal(model: pv.PolyData) -> float: + x_min, x_max, y_min, y_max, z_min, z_max = model.bounds + extents = np.array( + [x_max - x_min, y_max - y_min, z_max - z_min], dtype=np.float64 + ) + return float(np.linalg.norm(extents)) + + moving_diagonal = _diagonal(moving_model) + fixed_diagonal = _diagonal(fixed_model) + self.log_debug( + "Bounding-box diagonals - moving: %.4f, fixed: %.4f", + moving_diagonal, + fixed_diagonal, + ) + if moving_diagonal <= 0.0 or fixed_diagonal <= 0.0: + self.log_warning( + "Degenerate bounding box (moving diagonal %.4f, fixed diagonal " + "%.4f); skipping the scaling step.", + moving_diagonal, + fixed_diagonal, + ) + return 1.0 + return fixed_diagonal / moving_diagonal + + def _scale_transform(self, scale: float, center: np.ndarray) -> itk.AffineTransform: + """Build the transform scaling isotropically about ``center``. + + Args: + scale: Isotropic scale factor. + center: Fixed point of the scaling, in world coordinates. + + Returns: + An ITK affine point transform mapping ``p`` to + ``center + scale * (p - center)``. + """ + matrix = np.eye(3, dtype=np.float64) * scale + offset = itk.Vector[itk.D, 3]() + for i in range(3): + offset[i] = float(center[i]) * (1.0 - scale) + + transform = itk.AffineTransform[itk.D, 3].New() + transform.SetIdentity() + transform.SetMatrix(itk.Matrix[itk.D, 3, 3](itk.GetVnlMatrixFromArray(matrix))) + transform.SetOffset(offset) + return transform + def register( self, moving_model: pv.PolyData, @@ -141,22 +262,31 @@ def register( ) -> dict: """Perform ICP alignment of moving model to fixed model. - This method executes alignment with either rigid or affine transformations: + **Rigid transform type** (rotation + translation): + 1. Centroid alignment: Translate moving model to align mass centers + 2. Rigid ICP: Refine with rigid-body transformation - **Rigid transform type:** + **Similarity transform type** (rotation + translation + one uniform scale): 1. Centroid alignment: Translate moving model to align mass centers - 2. Rigid ICP: Refine with rigid-body transformation (rotation + translation) + 2. Bounding-box scaling: Scale isotropically about the fixed centroid + so the two bounding-box diagonals match, which keeps ICP's + closest-point search from locking onto a size mismatch + 3. Similarity ICP: Refine rotation, translation and uniform scale - **Affine transform type:** + **Affine transform type** (adds anisotropic scale and shear): 1. Centroid alignment: Translate moving model to align mass centers - 2. Rigid ICP: Refine with rigid-body transformation - 3. Affine ICP: Further refine with affine transformation (includes - scaling/shearing) + 2. Bounding-box scaling: As above + 3. Similarity ICP: Refine rotation, translation and uniform scale + 4. Affine ICP: Further refine with affine transformation + + 'Rigid' skips the bounding-box scaling so its result stays a pure + rigid-body transform; use 'Similarity' when the models differ in size but + the shape should not be distorted. Args: moving_model: PyVista surface model to be aligned to fixed model - transform_type: Registration transform type, either 'Rigid' or 'Affine'. - Default: 'Affine' + transform_type: Registration transform type, one of 'Rigid', + 'Similarity' or 'Affine'. Default: 'Affine' max_iterations: Maximum number of ICP iterations per stage. Default: 2000 Returns: @@ -168,7 +298,7 @@ def register( (ITK AffineTransform) Raises: - ValueError: If transform_type is not 'Rigid' or 'Affine' + ValueError: If transform_type is not 'Rigid', 'Similarity' or 'Affine' Example: >>> # Rigid registration @@ -178,6 +308,13 @@ def register( ... moving_model=moving_model, ... ) >>> + >>> # Similarity registration (rigid plus one uniform scale) + >>> result = registrar.register( + ... transform_type='Similarity', + ... max_iterations=2000, + ... moving_model=moving_model, + ... ) + >>> >>> # Affine registration >>> result = registrar.register( ... transform_type='Affine', @@ -185,9 +322,10 @@ def register( ... moving_model=moving_model, ... ) """ - if transform_type not in ["Rigid", "Affine"]: + if transform_type not in self._ICP_STAGES: raise ValueError( - f"Invalid transform '{transform_type}'. Must be 'Rigid' or 'Affine'." + f"Invalid transform '{transform_type}'. Must be one of " + f"{sorted(self._ICP_STAGES)}." ) self.log_section("%s ICP Alignment", transform_type.upper()) @@ -195,15 +333,15 @@ def register( self.moving_model = moving_model self.transform_type = transform_type - # Step 1: Centroid alignment (common to both modes) - self.registered_model = self.moving_model.copy(deep=True) + # Centroid alignment (common to every mode) + registered_model = self.moving_model.copy(deep=True) - moving_centroid = np.array(self.registered_model.center) + moving_centroid = np.array(registered_model.center) self.log_debug("Moving model centroid: %s", moving_centroid) fixed_centroid = np.array(self.fixed_model.center) self.log_debug("Fixed model centroid: %s", fixed_centroid) translation = fixed_centroid - moving_centroid - self.log_info("Step 1: Translating by %s to align centroids...", translation) + self.log_info("Translating by %s to align centroids...", translation) # Create ITK affine transform with translation forward_point_transform = itk.AffineTransform[itk.D, 3].New() @@ -211,70 +349,55 @@ def register( forward_point_transform.SetOffset(translation) # Apply centroid alignment to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, + registered_model = self.transform_tools.transform_pvcontour( + registered_model, forward_point_transform, with_deformation_magnitude=False, ) - self.log_debug("Center after Step 1: %s", self.registered_model.center) - - # Step 2: Rigid ICP (common to both modes) - self.log_info( - "Step 2: Performing rigid ICP (max iterations: %d)...", max_iterations - ) - icp_rigid = vtk.vtkIterativeClosestPointTransform() - icp_rigid.SetSource(self.registered_model) - icp_rigid.SetTarget(self.fixed_model) - icp_rigid.GetLandmarkTransform().SetModeToRigidBody() # Rigid mode - icp_rigid.SetMaximumNumberOfIterations(max_iterations) - icp_rigid.Update() - - # Convert VTK transform to ITK and compose with centroid transform - rigid_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( - icp_rigid.GetMatrix() - ) - forward_point_transform.Compose(rigid_transform) - - # Apply rigid ICP transform to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, - rigid_transform, - with_deformation_magnitude=False, - ) - - self.log_debug("Center after Step 2: %s", self.registered_model.center) - - # Step 3: Affine ICP (only if affine mode) - if transform_type == "Affine": + self.log_debug("Center after centroid alignment: %s", registered_model.center) + + # Bounding-box scaling, for the modes whose transform admits a scale. ICP + # only searches for correspondences among nearest points, so a template + # that differs from the patient in overall size drags the closest-point + # matching into a local minimum. Matching the bounding-box diagonals first + # puts the two models on the same scale before any ICP iteration runs. + # 'Rigid' skips it: a size estimate there would make the result a + # similarity transform, which is what 'Similarity' is for. + if transform_type != "Rigid": + scale = self._bounding_box_scale(registered_model, self.fixed_model) self.log_info( - "Step 3: Performing affine ICP (max iterations: %d)...", max_iterations - ) - icp_affine = vtk.vtkIterativeClosestPointTransform() - icp_affine.SetSource(self.registered_model) - icp_affine.SetTarget(self.fixed_model) - icp_affine.GetLandmarkTransform().SetModeToAffine() # Affine mode - icp_affine.SetMaximumNumberOfIterations(max_iterations) - icp_affine.Update() - - # Convert VTK transform to ITK and compose - affine_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( - icp_affine.GetMatrix() + "Scaling by %.4f about the fixed centroid to match bounding boxes...", + scale, ) - forward_point_transform.Compose(affine_transform) - - # Apply affine ICP transform to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, - affine_transform, + scale_transform = self._scale_transform(scale, fixed_centroid) + forward_point_transform.Compose(scale_transform) + registered_model = self.transform_tools.transform_pvcontour( + registered_model, + scale_transform, with_deformation_magnitude=False, ) + self.log_debug("Bounds after scaling: %s", registered_model.bounds) - self.log_debug("Center after Step 3: %s", self.registered_model.center) + # ICP stages. Affine runs similarity first so the shear and anisotropic + # scale degrees of freedom refine an already-oriented, already-sized model + # rather than absorbing rotation and overall scale themselves. + for stage in self._ICP_STAGES[transform_type]: + self.log_info( + "Performing %s ICP (max iterations: %d)...", + stage.lower(), + max_iterations, + ) + registered_model, stage_transform = self._icp_stage( + registered_model, stage, max_iterations + ) + forward_point_transform.Compose(stage_transform) + self.log_debug("Center after %s ICP: %s", stage, registered_model.center) # Compute inverse transform # Ths forward transform for ICP is consistent with the transform convention # used with images-to-images registration. + self.registered_model = registered_model self.forward_point_transform = forward_point_transform self.inverse_point_transform = forward_point_transform.GetInverseTransform() diff --git a/src/physiotwin4d/register_time_series_images.py b/src/physiotwin4d/register_time_series_images.py index 074fb320..3462588d 100644 --- a/src/physiotwin4d/register_time_series_images.py +++ b/src/physiotwin4d/register_time_series_images.py @@ -41,7 +41,6 @@ class RegisterTimeSeriesImages(RegisterImagesBase): - Sequential registration of ordered image lists - Supports any RegisterImagesBase backend, including RegisterImagesChain / RegisterImagesGreedyICON for multi-stage registration - - Optional use of prior transforms to initialize next registration - Configurable starting point in the time series - Returns all transforms and loss values for the entire series @@ -60,7 +59,6 @@ class RegisterTimeSeriesImages(RegisterImagesBase): ... moving_images=time_series_images, ... reference_frame=5, # Start from middle of cardiac cycle ... register_reference=True, - ... prior_weight=0.5, ... ) >>> >>> forward_tfms = result['forward_transforms'] # warp moving images -> fixed grid @@ -105,18 +103,6 @@ def __init__( self.transform_tools: TransformTools = TransformTools() - self.smooth_prior_transform_sigma: float = 0.5 - - def set_smooth_prior_transform_sigma( - self, smooth_prior_transform_sigma: float - ) -> None: - """Set the sigma for smoothing the prior transform. - - Args: - smooth_prior_transform_sigma (float): Sigma for smoothing the prior transform. - """ - self.smooth_prior_transform_sigma = smooth_prior_transform_sigma - def set_mask_dilation(self, mask_dilation_mm: float) -> None: """Set the dilation of the fixed and moving image masks. @@ -175,7 +161,6 @@ def register_time_series( moving_labelmaps: Optional[list[Optional[itk.Image]]] = None, reference_frame: int = 0, register_reference: bool = True, - prior_weight: float = 0.0, ) -> dict[str, list[itk.Transform] | list[float]]: """Register a time series of images to the fixed image. @@ -201,12 +186,7 @@ def register_time_series( register_reference (bool, optional): If True, register the reference image to the fixed image. If False, use identity transform for the reference image. Default: True - prior_weight (float, optional): - Weight (0.0 to 1.0) for using the prior image's transform to - initialize the next registration. 0.0 means no prior information - is used (each registration starts from identity). Higher values - provide more temporal smoothness but may propagate errors. - Default: 0.0 + Returns: dict: Dictionary containing results: - "forward_transforms" (list[itk.Transform]): one per image; @@ -222,13 +202,11 @@ def register_time_series( Raises: ValueError: If fixed_image is not set ValueError: If reference_frame is out of range - ValueError: If prior_weight not in [0, 1] ValueError: If moving_masks length doesn't match moving_images length Note: - The method compares registration with identity initialization versus - prior transform initialization and selects the result with lower loss. - This helps prevent error propagation in the temporal sequence. + Every frame is registered independently, so an error in one frame + cannot propagate along the series. The fixed image mask can be set using set_fixed_mask() before calling this method. @@ -245,7 +223,6 @@ def register_time_series( ... moving_labelmaps=labelmap_list, # Optional ... reference_frame=5, ... register_reference=True, - ... prior_weight=0.5, ... ) >>> >>> # Access results using new intuitive names @@ -273,9 +250,6 @@ def register_time_series( f"reference_frame {reference_frame} out of range [0, {num_images - 1}]" ) - if not 0.0 <= prior_weight <= 1.0: - raise ValueError("prior_weight must be in [0.0, 1.0]") - if moving_masks is not None and len(moving_masks) != num_images: raise ValueError( f"moving_masks length ({len(moving_masks)}) must match " @@ -329,29 +303,11 @@ def register_time_series( inverse_transforms[reference_frame] = inverse_transform losses[reference_frame] = loss - # Compute prior transform for reference frame if needed - prior_forward_ref = None - if prior_weight > 0.0: - prior_forward_ref = ( - self.transform_tools.combine_displacement_field_transforms( - identity_tfm, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=prior_weight, - tfm1_blur_sigma=0.0, - tfm2_blur_sigma=0.5, - mode="add", - ) - ) - # Register forward and backward from reference frame for step, start_idx, end_idx in [ (1, reference_frame + 1, num_images), # Forward pass (-1, reference_frame - 1, -1), # Backward pass ]: - prior_forward = prior_forward_ref - for img_idx in range(start_idx, end_idx, step): moving_image = moving_images[img_idx] moving_mask = ( @@ -361,65 +317,15 @@ def register_time_series( moving_labelmaps[img_idx] if moving_labelmaps is not None else None ) - # Try registration with identity initialization - result_init_identity = self.registrar.register( + result = self.registrar.register( moving_image=moving_image, moving_mask=moving_mask, moving_labelmap=moving_labelmap, ) - forward_init_identity = result_init_identity["forward_transform"] - inverse_init_identity = result_init_identity["inverse_transform"] - loss_init_identity = result_init_identity["loss"] - - # Select best result based on prior usage - if prior_weight > 0.0: - # Try with prior transform initialization - result_init_prior = self.registrar.register( - moving_image=moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, - initial_forward_transform=prior_forward, - ) - forward_init_prior = result_init_prior["forward_transform"] - inverse_init_prior = result_init_prior["inverse_transform"] - loss_init_prior = result_init_prior["loss"] - - # Select result with lower loss - if loss_init_identity < loss_init_prior: - # Identity initialization was better - prior_forward = identity_tfm - forward_transform = forward_init_identity - inverse_transform = inverse_init_identity - loss = loss_init_identity - else: - # Prior initialization was better - forward_transform = forward_init_prior - inverse_transform = inverse_init_prior - loss = loss_init_prior - - # Update prior for next iteration - prior_forward = ( - self.transform_tools.combine_displacement_field_transforms( - identity_tfm, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=prior_weight, - tfm1_blur_sigma=0.0, - tfm2_blur_sigma=self.smooth_prior_transform_sigma, - mode="add", - ) - ) - else: - # No prior usage, just use identity result - forward_transform = forward_init_identity - inverse_transform = inverse_init_identity - loss = loss_init_identity - - # Store results - forward_transforms[img_idx] = forward_transform - inverse_transforms[img_idx] = inverse_transform - losses[img_idx] = loss + + forward_transforms[img_idx] = result["forward_transform"] + inverse_transforms[img_idx] = result["inverse_transform"] + losses[img_idx] = cast(float, result["loss"]) assert all(t is not None for t in forward_transforms) assert all(t is not None for t in inverse_transforms) @@ -567,7 +473,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Registration method required by RegisterImagesBase. @@ -581,7 +486,6 @@ def registration_method( moving_labelmap (itk.Image, optional): Multi-label segmentation moving_image_pre (itk.Image, optional): Ignored - the registrar computes its own preprocessing from the raw moving_image - initial_forward_transform (itk.Transform, optional): Initial transform Returns: dict: Registration result with forward_transform, inverse_transform, and loss @@ -592,7 +496,6 @@ def registration_method( moving_mask=moving_mask, moving_labelmap=moving_labelmap, moving_image_pre=None, - initial_forward_transform=initial_forward_transform, ) self._capture_delegate_result(self.registrar, result) return { diff --git a/src/physiotwin4d/train_physicsnemo_mgn.py b/src/physiotwin4d/train_physicsnemo_mgn.py index f50dcd74..82338eb4 100644 --- a/src/physiotwin4d/train_physicsnemo_mgn.py +++ b/src/physiotwin4d/train_physicsnemo_mgn.py @@ -66,6 +66,22 @@ def set_num_layers(self, num_layers: int) -> None: raise ValueError(f"num_layers must be >= 1, got {num_layers}") self.num_layers = num_layers + def set_num_processor_checkpoint_segments(self, num_segments: int) -> None: + """Set the gradient-checkpointing segment count for the processor. + + Gradient checkpointing recomputes processor activations during the + backward pass instead of storing them, trading compute for GPU memory. + A large mesh graph makes that trade worthwhile: a 179k-point, 1.07M-edge + template peaks near 43 GiB at ``batch_size`` 4 without it. + + Args: + num_segments: Number of checkpointed segments; ``0`` (the default) + disables checkpointing and stores every activation. + """ + if num_segments < 0: + raise ValueError(f"num_segments must be >= 0, got {num_segments}") + self.num_processor_checkpoint_segments = num_segments + def build_model(self, in_features: int, out_features: int) -> "torch.nn.Module": try: import torch_geometric # noqa: F401 - needed by the graph seams diff --git a/src/physiotwin4d/transform_tools.py b/src/physiotwin4d/transform_tools.py index 75f86efd..d6e119fb 100644 --- a/src/physiotwin4d/transform_tools.py +++ b/src/physiotwin4d/transform_tools.py @@ -255,9 +255,30 @@ def convert_transform_to_displacement_field_transform( new_tfm.SetDisplacementField(field) return new_tfm - def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transform: + def invert_displacement_field_transform( + self, + tfm: itk.Transform, + max_iterations: int = 20, + max_error_tolerance: float = 0.05, + mean_error_tolerance: float = 0.0005, + ) -> itk.Transform: """ Invert a displacement field transform. + + Uses SimpleITK's fixed-point iterative inversion on the input field's + own grid. The defaults are tighter than SimpleITK's own (10 iterations, + 0.1 mm max error) because the fields produced here are not smooth + everywhere and converge slowly near their support boundary. + + Args: + tfm: Displacement field transform to invert. + max_iterations: Fixed-point iterations per voxel. + max_error_tolerance: Convergence threshold on the maximum error, in + the field's units. + mean_error_tolerance: Convergence threshold on the mean error. + + Returns: + The inverted displacement field transform. """ assert "DisplacementFieldTransform" in str(type(tfm)), ( "Input transform must be a displacement field transform" @@ -268,7 +289,12 @@ def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transfo field_sitk = image_tools.convert_itk_image_to_sitk(field_itk) - field_sitk_inv = sitk.InvertDisplacementField(field_sitk) + field_sitk_inv = sitk.InvertDisplacementField( + field_sitk, + maximumNumberOfIterations=max_iterations, + maxErrorToleranceThreshold=max_error_tolerance, + meanErrorToleranceThreshold=mean_error_tolerance, + ) field_itk_inv = image_tools.convert_sitk_image_to_itk(field_sitk_inv) @@ -277,6 +303,38 @@ def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transfo return new_tfm + def invert_transform( + self, tfm: itk.Transform, reference_image: itk.Image + ) -> itk.Transform: + """Invert any transform, analytically when the type supports it. + + Prefers ITK's analytic inverse (available for translation, rigid, affine + and composites of them) and falls back to rasterizing a displacement + field over ``reference_image`` and inverting that numerically. + + The analytic inverse is preferred because the fallback is only defined on + ``reference_image``'s grid: outside it the field is zero, so the inverse + silently degrades to the identity there. + + Args: + tfm (itk.Transform): Transform to invert. + reference_image (itk.Image): Grid used by the displacement-field + fallback. + + Returns: + itk.Transform: The inverse transform. + """ + try: + analytic = tfm.GetInverseTransform() + except Exception: # pragma: no cover - transform type dependent + analytic = None + if analytic is not None: + return cast(itk.Transform, analytic) + + return self.invert_displacement_field_transform( + self.convert_transform_to_displacement_field_transform(tfm, reference_image) + ) + def transform_pvcontour( self, contour: pv.PolyData, diff --git a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py index 44041aaf..4b15ec06 100644 --- a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py @@ -480,12 +480,17 @@ def register_model_to_model_icp(self) -> dict: def register_model_to_model_pca(self) -> dict: """Perform PCA-based registration after ICP alignment. - Uses RegisterModelsPCA class for intensity-based PCA registration. - This method requires PCA data to be set via set_pca_data(). + Uses RegisterModelsPCA to optimize shape coefficients against a distance + map of the patient. The statistical model's modes are defined in the + un-aligned template frame, so the registrar is given the raw template + and the ICP alignment is passed as its ``post_pca_transform``. Returns: dict: Dictionary containing: - - 'forward_point_transform': Rigid transform from PCA registration + - 'forward_point_transform': DisplacementFieldTransform mapping + un-aligned template points to their PCA-deformed positions. + It excludes the ICP alignment, which is applied separately. + - 'inverse_point_transform': its inverse - 'pca_coefficients': PCA shape coefficients - 'registered_template_model_surface': PCA-registered model surface @@ -514,13 +519,18 @@ def register_model_to_model_pca(self) -> dict: "inverse_point_transform": self.pca_inverse_point_transform, } + # PCA modes are directions in the statistical model's own training + # frame, so they must be added to the un-aligned template and the ICP + # alignment applied afterwards. Deforming the ICP-aligned template + # instead would yield A*mean + sum(b*sigma*v) rather than + # A*(mean + sum(b*sigma*v)), mis-rotating and mis-scaling every mode. pca_template_model: Optional[pv.DataSet] if self.use_surface: - pca_template_model = self.icp_template_model_surface + pca_template_model = self.template_model_surface fixed_model = self.patient_model_surface fixed_distance_map = None else: - pca_template_model = self.icp_template_model + pca_template_model = self.template_model fixed_model = self.combined_patient_model if self.patient_labelmap is not None: fixed_distance_map = self.labelmap_tools.create_distance_map( @@ -539,6 +549,7 @@ def register_model_to_model_pca(self) -> dict: pca_template_model=pca_template_model, pca_model=self.pca_model, pca_number_of_modes=self.pca_number_of_modes, + post_pca_transform=self.icp_forward_point_transform, fixed_model=fixed_model, fixed_distance_map=fixed_distance_map, reference_image=self.patient_image, @@ -580,11 +591,20 @@ def register_model_to_model_pca(self) -> dict: itk.imwrite(tfm_z_img, "pca_forward_point_transform_z.nii.gz") if self.use_surface: - assert self.icp_template_model is not None, "ICP template model must be set" - self.pca_template_model = self._transform_model_dataset( - self.icp_template_model, + # forward_point_transform excludes the post-PCA step and is defined + # in the un-aligned template frame, so warp the raw volumetric + # template with it and then apply the ICP alignment. + assert self.icp_forward_point_transform is not None, ( + "ICP forward transform must be set" + ) + deformed_template_model = self._transform_model_dataset( + self.template_model, self.pca_forward_point_transform, ) + self.pca_template_model = self._transform_model_dataset( + deformed_template_model, + self.icp_forward_point_transform, + ) else: self.pca_template_model = registered_model @@ -593,10 +613,22 @@ def register_model_to_model_pca(self) -> dict: self.registered_template_model = self.pca_template_model self.registered_template_model_surface = self.pca_template_model_surface - if self.icp_template_labelmap is not None: + if self.template_labelmap is not None: + # Resampling pulls back: a patient-grid sample is mapped by the ICP + # inverse into the deformed-template frame and then by the PCA + # inverse onto the un-deformed template. itk.CompositeTransform + # applies its transforms in reverse order of addition, so the ICP + # inverse is added last. Resampling the raw template labelmap in one + # step also avoids a second round of nearest-neighbor sampling. + assert self.icp_inverse_point_transform is not None, ( + "ICP inverse transform must be set" + ) + pca_image_transform = itk.CompositeTransform[itk.D, 3].New() + pca_image_transform.AddTransform(self.pca_inverse_point_transform) + pca_image_transform.AddTransform(self.icp_inverse_point_transform) self.pca_template_labelmap = self.transform_tools.transform_image( - self.icp_template_labelmap, - self.pca_inverse_point_transform, + self.template_labelmap, + pca_image_transform, self.patient_image, interpolation_method="nearest", ) @@ -642,11 +674,19 @@ def register_labelmap_to_labelmap(self) -> Optional[dict]: assert self.pca_template_model_surface is not None, ( "PCA template model surface must be set" ) + + # Create a padded patient image since often the surface of interest + # is not fully contained within the original image, which causes trouble + # with distance map registration. + padded_patient_image = ImageTools().pad_image( + self.patient_image, pad_voxels=[50, 50, 50], background_value=-1000 + ) labelmap_registrar = RegisterModelsDistanceMaps( moving_model=self.pca_template_model_surface, fixed_model=self.patient_model_surface, - reference_image=self.patient_image, + reference_image=padded_patient_image, mask_dilation_mm=self.mask_dilation_mm, + distance_squared_max=(1.25 * self.mask_dilation_mm) ** 2, ) # Run deformable registration @@ -772,10 +812,10 @@ def register_labelmap_to_image( self.registrar_ICON.set_fixed_image(self.patient_image) self.registrar_ICON.set_fixed_mask(patient_mask) - # Perform Icon registration - result = self.registrar_ICON.register( - initial_forward_transform=self.l2i_forward_transform, - moving_image=template_labelmap, + # Perform Icon registration, refining the alignment found so far + result = self.registrar_ICON.register_from( + self.l2i_forward_transform, + template_labelmap, moving_mask=template_mask, ) self.l2i_inverse_transform = result["inverse_transform"] @@ -847,9 +887,10 @@ def transform_model( transformed_model = base_model.copy(deep=True) transform_steps: list[tuple[str, itk.Transform]] = [] - if self.icp_forward_point_transform is not None: - transform_steps.append(("ICP", self.icp_forward_point_transform)) if self.pca_coefficients is not None: + # PCA registration runs in the un-aligned template frame and carries + # the ICP alignment in its post-PCA transform, so ICP must not be + # applied again here. assert self.pca_registrar is not None, "PCA registrar must be set" pca_transform = ( self.pca_forward_point_transform @@ -861,6 +902,8 @@ def transform_model( transform_steps.append( ("PCA post-transform", self.pca_registrar.post_pca_transform) ) + elif self.icp_forward_point_transform is not None: + transform_steps.append(("ICP", self.icp_forward_point_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: @@ -908,9 +951,7 @@ def process( Returns: dict with registered_template_model and registered_template_model_surface """ - self.log_section( - "STARTING COMPLETE MODEL-TO-IMAGE-AND-MODEL REGISTRATION WORKFLOW", width=70 - ) + self.log_section("STARTING COMPLETE MODEL REGISTRATION WORKFLOW", width=70) self.use_ICON_registration_refinement = use_ICON_registration_refinement diff --git a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py index a6c96951..3d1ed2db 100644 --- a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py @@ -64,7 +64,6 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): reference_time_frame (int): Index of reference time frame in time series register_reference_time_frame_to_reference_image (bool): Whether to register the reference time frame to the reference image - prior_weight (float): Weight for temporal smoothing (0.0-1.0) upsample_to_fixed_resolution (bool): Whether to upsample reconstruction registrar (RegisterTimeSeriesImages): Internal registration object forward_transforms (list[itk.Transform]): one per frame; each warps its @@ -152,7 +151,6 @@ def __init__( ) # Initialize parameters with defaults - self.prior_weight: float = 0.0 self.upsample_to_fixed_resolution: bool = True self.modality: str = "ct" self.mask_dilation_mm: float = 0.0 @@ -170,22 +168,6 @@ def __init__( self.losses: Optional[list[float]] = None self.reconstructed_images: Optional[list[itk.Image]] = None - def set_prior_weight(self, prior_weight: float) -> None: - """Set the weight for temporal smoothing with prior transforms. - - Args: - prior_weight (float): Weight (0.0 to 1.0) for using the prior image's - transform to initialize the next registration. 0.0 means no prior - information is used (each registration starts from identity). - Higher values provide more temporal smoothness but may propagate errors. - - Raises: - ValueError: If prior_weight not in [0.0, 1.0] - """ - if not 0.0 <= prior_weight <= 1.0: - raise ValueError(f"prior_weight must be in [0.0, 1.0], got {prior_weight}") - self.prior_weight = prior_weight - def set_modality(self, modality: str) -> None: """Set the imaging modality for registration optimization. @@ -267,7 +249,6 @@ def register_time_series(self) -> dict: "Register reference time frame to reference image: " f"{self.register_reference_time_frame_to_reference_image}" ) - self.log_info(f"Prior weight: {self.prior_weight}") # Perform registration result = self.registrar.register_time_series( @@ -275,7 +256,6 @@ def register_time_series(self) -> dict: moving_masks=self.moving_masks, reference_frame=self.reference_time_frame, register_reference=self.register_reference_time_frame_to_reference_image, - prior_weight=self.prior_weight, ) # Store results @@ -378,7 +358,6 @@ def process(self) -> dict: registrar_type = type(self.registrar.registrar).__name__ self.log_info(f" Registration method: {registrar_type}") self.log_info(f" Reference time frame: {self.reference_time_frame}") - self.log_info(f" Prior weight: {self.prior_weight}") self.log_info( f" Upsample to fixed resolution: {self.upsample_to_fixed_resolution}" ) diff --git a/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf index 75823466..229d4d4f 100644 --- a/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf +++ b/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:24697b25fa57cc7d83ad221b42789e75409ef96470a1ad2d89981b90386ccfc6 -size 40878719 +oid sha256:3fdb6bab6c362e6d8b2f7f23db540e24eb6b732e7f6cb82ce713433daa0670bd +size 4512579 diff --git a/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha b/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha index 3098be93..73f512a1 100644 --- a/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha +++ b/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 +oid sha256:fa8492f316a3c9c78b75a06ef45152fc8f5a8db36b0931386ef1ca22a6f6197e +size 4767196 diff --git a/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf new file mode 100644 index 00000000..afb51242 --- /dev/null +++ b/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b6d8ed46f2654a646111b6c6aad58641bf35573277d339d8d02d5e27e6737c4 +size 4158463 diff --git a/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf deleted file mode 100644 index 664096cd..00000000 --- a/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fd4bb595e748dcc777ffe290b18f8b1a6436bc2e579b298bc529157326e7b40f -size 40878719 diff --git a/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha b/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha deleted file mode 100644 index 3098be93..00000000 --- a/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 diff --git a/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha b/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha index 3098be93..e2ec568b 100644 --- a/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha +++ b/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 +oid sha256:5615af55c808458fc9e9e7c999ad4ddfc85a6082f2db1bfb6ec5bc770c03d892 +size 4779024 diff --git a/tests/conftest.py b/tests/conftest.py index 31b5a686..0caf99aa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,6 +12,7 @@ from typing import Any, Optional import itk +import numpy as np import pytest from physiotwin4d.contour_tools import ContourTools @@ -662,3 +663,75 @@ def registrar_ICON() -> RegisterImagesICON: def transform_tools() -> TransformTools: """Create a TransformTools instance.""" return TransformTools() + + +class KnownShiftCase: + """A registration case whose correct answer is known exactly. + + ``moving`` is built by resampling ``fixed`` through a translation of + ``shift_mm``, so ``moving(q) == fixed(q + shift_mm)``. Warping ``moving`` + back onto the fixed grid therefore requires a ``forward_transform`` of + ``-shift_mm``, which gives an absolute accuracy target instead of the + "did it return something" checks that let a Greedy RAS/LPS sign error go + unnoticed. + """ + + def __init__(self, fixed_image: itk.Image, shift_mm: tuple[float, float, float]): + """Build the shifted pair. + + Args: + fixed_image: Image used as the registration target. + shift_mm: Content displacement applied to build the moving image. + Use a different magnitude and sign per axis so an axis swap or a + sign flip cannot pass. + """ + self.transform_tools = TransformTools() + self.fixed = fixed_image + self.shift_mm = shift_mm + self.expected_displacement = np.array([-v for v in shift_mm]) + + shift = itk.TranslationTransform[itk.D, 3].New() + shift.SetOffset(list(shift_mm)) + self.moving = self.transform_tools.transform_image( + fixed_image, shift, fixed_image, interpolation_method="linear" + ) + + size = itk.size(fixed_image) + self._center = list( + fixed_image.TransformIndexToPhysicalPoint( + [int(size[i]) // 2 for i in range(3)] + ) + ) + # Score over the brightest 30% of the fixed image (tissue and blood + # pool); background air correlates trivially and would mask errors. + self._fixed_array = itk.array_from_image(fixed_image) + self._foreground = self._fixed_array >= np.percentile(self._fixed_array, 70) + + def center_error_mm(self, forward_transform: itk.Transform) -> float: + """Distance, in mm, between the recovered and true displacement.""" + displacement = np.array( + list(forward_transform.TransformPoint(self._center)) + ) - np.array(self._center) + return float(np.linalg.norm(displacement - self.expected_displacement)) + + def foreground_ncc(self, forward_transform: itk.Transform) -> float: + """Normalized cross-correlation after warping moving onto the fixed grid.""" + warped = self.transform_tools.transform_image( + self.moving, forward_transform, self.fixed, interpolation_method="linear" + ) + moved = itk.array_from_image(warped)[self._foreground] + target = self._fixed_array[self._foreground] + moved = moved - moved.mean() + target = target - target.mean() + denominator = np.sqrt((moved**2).sum() * (target**2).sum()) + return float((moved * target).sum() / denominator) if denominator else 0.0 + + def unregistered_ncc(self) -> float: + """Baseline score with no registration, for a floor to beat.""" + return self.foreground_ncc(itk.TranslationTransform[itk.D, 3].New()) + + +@pytest.fixture(scope="session") +def known_shift_case(test_images: list[Any]) -> KnownShiftCase: + """A moving/fixed pair separated by a known (6, -4, 3) mm shift.""" + return KnownShiftCase(test_images[0], (6.0, -4.0, 3.0)) diff --git a/tests/test_register_images_ants.py b/tests/test_register_images_ants.py index a40a5237..862512cb 100644 --- a/tests/test_register_images_ants.py +++ b/tests/test_register_images_ants.py @@ -18,6 +18,8 @@ from physiotwin4d.register_images_ants import RegisterImagesANTS from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + def _foreground_ncc( reference_arr: np.ndarray, warped_arr: np.ndarray, foreground: np.ndarray @@ -294,6 +296,35 @@ def test_preprocess_images( print("Image preprocessing complete") print(f" Preprocessed spacing: {preprocessed_spacing}") + def test_recovers_known_shift(self, known_shift_case: KnownShiftCase) -> None: + """ANTs must recover a known shift, in the right direction. + + The companion of the Greedy check in test_register_images_greedy.py: + ANTs works in ITK's LPS frame throughout, so this asserts no equivalent + RAS/LPS conversion is missing here. + """ + registrar = RegisterImagesANTS() + registrar.set_modality("ct") + registrar.set_transform_type("Affine") + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print("\nANTs known-shift recovery:") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, f"ANTs recovered the shift {error_mm:.2f} mm off" + assert ncc > unregistered_ncc, ( + f"ANTs left the images less aligned than they started " + f"({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_registration_with_initial_transform( self, registrar_ANTS: RegisterImagesANTS, @@ -318,9 +349,9 @@ def test_registration_with_initial_transform( registrar_ANTS.set_modality("ct") registrar_ANTS.set_fixed_image(fixed_image) - result = registrar_ANTS.register( - moving_image=moving_image, - initial_forward_transform=initial_tfm_forward, + result = registrar_ANTS.register_from( + initial_tfm_forward, + moving_image, ) assert isinstance(result, dict), "Result should be a dictionary" @@ -335,7 +366,7 @@ def test_initial_transform_composition_metrics( test_images: list[Any], test_directories: dict[str, Path], ) -> None: - """Verify the initial_forward_transform composition path with metrics. + """Verify the register_from() composition path with metrics. Exercises the two initial-transform inputs the platform actually uses (identity and a prior deformable forward_transform, as in prior-based @@ -396,18 +427,16 @@ def warp_score(forward_transform: Any) -> float: registrar_identity = RegisterImagesANTS() registrar_identity.set_modality("ct") registrar_identity.set_fixed_image(fixed_image) - identity_result = registrar_identity.register( - moving_image=moving_image, initial_forward_transform=identity - ) + identity_result = registrar_identity.register_from(identity, moving_image) ncc_identity = warp_score(identity_result["forward_transform"]) # Prior deformable initial: the realistic time-series prior use case. registrar_prior = RegisterImagesANTS() registrar_prior.set_modality("ct") registrar_prior.set_fixed_image(fixed_image) - prior_result = registrar_prior.register( - moving_image=moving_image, - initial_forward_transform=baseline["forward_transform"], + prior_result = registrar_prior.register_from( + baseline["forward_transform"], + moving_image, ) ncc_prior = warp_score(prior_result["forward_transform"]) @@ -475,9 +504,7 @@ def test_initial_transform_matrix_composition( registrar_ANTS.set_modality("ct") registrar_ANTS.set_fixed_image(fixed_image) - result = registrar_ANTS.register( - moving_image=moving_image, initial_forward_transform=translation - ) + result = registrar_ANTS.register_from(translation, moving_image) transform_tools = TransformTools() warped = transform_tools.transform_image( diff --git a/tests/test_register_images_chain.py b/tests/test_register_images_chain.py index 26e999c7..62665dfe 100644 --- a/tests/test_register_images_chain.py +++ b/tests/test_register_images_chain.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Union +from typing import Any, Optional, Union, cast import itk import numpy as np @@ -41,7 +41,6 @@ def __init__(self, name: str, sentinel_value: float) -> None: self.seen_moving_image: Optional[itk.Image] = None self.seen_moving_image_pre: Optional[itk.Image] = None self.seen_moving_mask: Optional[itk.Image] = None - self.seen_initial_forward_transform: Optional[object] = None self.preprocess_call_count = 0 def preprocess(self, image: itk.Image, modality: str = "ct") -> itk.Image: @@ -55,34 +54,47 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[object] = None, ) -> dict[str, Union[object, float]]: - """Record the state visible at call time; return a sentinel result.""" + """Record the state visible at call time; return a sentinel result. + + The sentinel transforms are real translations along x of + ``sentinel_value`` mm, so a chained result can be checked numerically. + """ self.seen_fixed_image_pre = self.fixed_image_pre self.seen_moving_image = self.moving_image self.seen_moving_image_pre = moving_image_pre self.seen_moving_mask = moving_mask - self.seen_initial_forward_transform = initial_forward_transform + forward = itk.TranslationTransform[itk.D, 3].New() + forward.SetOffset([self.sentinel_value, 0.0, 0.0]) + inverse = itk.TranslationTransform[itk.D, 3].New() + inverse.SetOffset([-self.sentinel_value, 0.0, 0.0]) return { - "forward_transform": f"forward_{self.name}", - "inverse_transform": f"inverse_{self.name}", + "forward_transform": forward, + "inverse_transform": inverse, "loss": self.sentinel_value, } -def test_chain_feeds_forward_transform_to_next_stage() -> None: - """Stage 2 must receive stage 1's forward_transform as its - initial_forward_transform.""" +def test_chain_refines_previous_stage_result() -> None: + """Stage 2 must refine stage 1's alignment rather than start over. + + Stage 1 is called on the raw moving image; stage 2 sees a moving image + pre-warped by stage 1's forward transform, and the chain's composed + forward transform is the sum of both stages' translations. + """ stage1 = _RecordingRegistrar("stage1", 1.0) stage2 = _RecordingRegistrar("stage2", 2.0) chain = RegisterImagesChain([stage1, stage2]) chain.set_fixed_image(_small_image()) + moving = _small_image() - result = chain.register(_small_image()) + result = chain.register(moving) - assert stage1.seen_initial_forward_transform is None - assert stage2.seen_initial_forward_transform == "forward_stage1" - assert result["forward_transform"] == "forward_stage2" + assert stage1.seen_moving_image is moving + # Stage 2 registers the pre-warped image, not the caller's image. + assert stage2.seen_moving_image is not moving + composed = cast(itk.Transform, result["forward_transform"]) + assert list(composed.TransformPoint([0.0, 0.0, 0.0])) == [3.0, 0.0, 0.0] assert result["loss"] == 2.0 @@ -101,12 +113,16 @@ def test_chain_propagates_fixed_and_moving_state_to_each_child() -> None: for stage in (stage1, stage2): assert stage.seen_fixed_image_pre is not None - assert stage.seen_moving_image is moving + assert stage.seen_moving_image is not None # Each stage must compute its own preprocessing (moving_image_pre is # not inherited from the chain, which has no meaningful preprocess() # of its own). assert stage.seen_moving_image_pre is None assert stage.preprocess_call_count == 1 + # The first stage gets the caller's image; later stages get it pre-warped + # by the running result. + assert stage1.seen_moving_image is moving + assert stage2.seen_moving_image is not moving def test_chain_recomputes_fixed_image_pre_when_fixed_image_changes() -> None: diff --git a/tests/test_register_images_greedy.py b/tests/test_register_images_greedy.py index d079c709..efdbf065 100644 --- a/tests/test_register_images_greedy.py +++ b/tests/test_register_images_greedy.py @@ -15,6 +15,8 @@ from physiotwin4d.register_images_greedy import RegisterImagesGreedy from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + @pytest.mark.slow class TestRegisterImagesGreedy: @@ -183,6 +185,48 @@ def test_register_affine_with_mask( print("Greedy affine registration complete with masks") + @pytest.mark.parametrize("transform_type", ["Rigid", "Affine", "Deformable"]) + def test_recovers_known_shift( + self, + known_shift_case: KnownShiftCase, + transform_type: str, + ) -> None: + """Greedy must recover a known shift, in the right direction. + + Regression guard for the RAS/LPS conversion: Greedy reports its affine + in RAS while ITK is LPS, so omitting the basis change negates x and y. + Before the fix this recovered (+6, -4, -3) mm instead of (-6, +4, -3) + and scored *below* the unregistered pair; the sign error passed every + other test in this file, which only check that transforms exist. + """ + registrar = RegisterImagesGreedy() + registrar.set_modality("ct") + registrar.set_transform_type(transform_type) + registrar.set_number_of_iterations([60, 30, 10]) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print(f"\nGreedy {transform_type} known-shift recovery:") + print(f" expected displacement: {known_shift_case.expected_displacement}") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, ( + f"Greedy {transform_type} recovered the shift {error_mm:.2f} mm off; " + "a sign or axis error inverts the transform (see the RAS/LPS " + "conversion in RegisterImagesGreedy._matrix_to_itk_affine)" + ) + assert ncc > unregistered_ncc, ( + f"Greedy {transform_type} left the images less aligned than they " + f"started ({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_transform_application( self, registrar_greedy: RegisterImagesGreedy, diff --git a/tests/test_register_images_icon.py b/tests/test_register_images_icon.py index 7ecde324..f3e369aa 100644 --- a/tests/test_register_images_icon.py +++ b/tests/test_register_images_icon.py @@ -16,12 +16,43 @@ from physiotwin4d.register_images_icon import RegisterImagesICON from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + @pytest.mark.requires_gpu @pytest.mark.slow class TestRegisterImagesICON: """Test suite for ICON-based image registration.""" + def test_recovers_known_shift(self, known_shift_case: KnownShiftCase) -> None: + """ICON must recover a known shift, in the right direction. + + The companion of the Greedy check in test_register_images_greedy.py: + ICON returns ITK transforms in the LPS frame throughout, so this asserts + no equivalent RAS/LPS conversion is missing here. + """ + registrar = RegisterImagesICON() + registrar.set_modality("ct") + registrar.set_number_of_iterations(5) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print("\nICON known-shift recovery:") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, f"ICON recovered the shift {error_mm:.2f} mm off" + assert ncc > unregistered_ncc, ( + f"ICON left the images less aligned than they started " + f"({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_registrar_initialization(self, registrar_ICON: RegisterImagesICON) -> None: """Test that RegisterImagesICON initializes correctly.""" assert registrar_ICON is not None, "Registrar not initialized" @@ -386,9 +417,9 @@ def test_registration_with_initial_transform( registrar_ICON.set_fixed_image(fixed_image) registrar_ICON.set_number_of_iterations(2) - result = registrar_ICON.register( - moving_image=moving_image, - initial_forward_transform=initial_tfm_forward, + result = registrar_ICON.register_from( + initial_tfm_forward, + moving_image, ) assert isinstance(result, dict), "Result should be a dictionary" diff --git a/tests/test_register_time_series_images.py b/tests/test_register_time_series_images.py index e7268d9b..10396c4b 100644 --- a/tests/test_register_time_series_images.py +++ b/tests/test_register_time_series_images.py @@ -145,7 +145,6 @@ def test_register_time_series_basic( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) # Verify result structure @@ -207,16 +206,15 @@ def test_register_time_series_basic( assert (results_dir / "basic_forward_transform_0.hdf").exists() assert (results_dir / "basic_time_series_registered_0.mha").exists() - def test_register_time_series_with_prior( + def test_register_time_series_from_middle_frame( self, test_images: list[Any], test_directories: dict[str, Path] ) -> None: - """Test time series registration with prior transform usage.""" + """Test time series registration starting from a middle reference frame.""" fixed_image = test_images[0] moving_images = test_images[1:4] - print("\nRegistering time series (with prior)...") + print("\nRegistering time series (middle reference frame)...") print(f" Number of moving images: {len(moving_images)}") - print(" Using prior transform weight: 0.5") greedy = RegisterImagesGreedy() greedy.set_number_of_iterations([20, 10, 2]) @@ -228,7 +226,6 @@ def test_register_time_series_with_prior( moving_images=moving_images, reference_frame=1, # Start from middle register_reference=True, - prior_weight=0.5, ) forward_transforms = result["forward_transforms"] @@ -246,7 +243,7 @@ def test_register_time_series_with_prior( for i, forward_transform in enumerate(forward_transforms): assert forward_transform is not None, f"forward_transform[{i}] is None" - print("Time series registration with prior complete") + print("Time series registration from the middle frame complete") print(f" Losses: {[f'{loss:.6f}' for loss in losses]}") test_tools = TestTools( @@ -259,14 +256,14 @@ def test_register_time_series_with_prior( # bit-reproducible across runs, so we save artifacts without # asserting an exact baseline match. test_tools.write_result_transform( - forward_transforms[0], "prior_forward_transform_0.hdf" + forward_transforms[0], "middle_frame_forward_transform_0.hdf" ) test_tools.write_result_image( - moving_image, "prior_time_series_registered_0.mha" + moving_image, "middle_frame_time_series_registered_0.mha" ) results_dir = test_directories["output"] / self._class_name - assert (results_dir / "prior_forward_transform_0.hdf").exists() - assert (results_dir / "prior_time_series_registered_0.mha").exists() + assert (results_dir / "middle_frame_forward_transform_0.hdf").exists() + assert (results_dir / "middle_frame_time_series_registered_0.mha").exists() def test_register_time_series_identity_start(self, test_images: list[Any]) -> None: """Test time series registration with identity for starting image.""" @@ -285,7 +282,6 @@ def test_register_time_series_identity_start(self, test_images: list[Any]) -> No moving_images=moving_images, reference_frame=0, register_reference=False, # Use identity - prior_weight=0.0, ) # Starting image should have very low/zero loss @@ -317,7 +313,6 @@ def test_register_time_series_different_starting_indices( moving_images=moving_images, reference_frame=starting_index, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images), ( @@ -360,31 +355,6 @@ def test_register_time_series_error_invalid_starting_index( print("\nInvalid starting index correctly rejected") - def test_register_time_series_error_invalid_prior_portion( - self, test_images: list[Any] - ) -> None: - """Test that error is raised for invalid prior portion value.""" - registrar = RegisterTimeSeriesImages(registration_method=RegisterImagesGreedy()) - registrar.set_fixed_image(test_images[0]) - - moving_images = test_images[1:4] - - # Test negative value - with pytest.raises(ValueError, match="must be in"): - registrar.register_time_series( - moving_images=moving_images, - prior_weight=-0.1, - ) - - # Test value > 1 - with pytest.raises(ValueError, match="must be in"): - registrar.register_time_series( - moving_images=moving_images, - prior_weight=1.5, - ) - - print("\nInvalid prior portion correctly rejected") - def test_transform_application_time_series( self, test_images: list[Any], test_directories: dict[str, Path] ) -> None: @@ -404,7 +374,6 @@ def test_transform_application_time_series( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) forward_transforms = result["forward_transforms"] @@ -457,7 +426,6 @@ def test_register_time_series_ICON(self, test_images: list[Any]) -> None: moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images) @@ -505,7 +473,6 @@ def test_register_time_series_with_mask( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images) @@ -531,7 +498,6 @@ def test_bidirectional_registration(self, test_images: list[Any]) -> None: moving_images=moving_images, reference_frame=2, # Middle image register_reference=True, - prior_weight=0.0, ) forward_transforms = result["forward_transforms"] diff --git a/tutorials/tutorial_06_lung_create_statistical_model.py b/tutorials/tutorial_06_lung_create_statistical_model.py index 02fcf5df..6890491a 100644 --- a/tutorials/tutorial_06_lung_create_statistical_model.py +++ b/tutorials/tutorial_06_lung_create_statistical_model.py @@ -67,7 +67,9 @@ baselines_dir = repo_root / "tests" / "baselines" data_dir = repo_root / "data" / "DirLab-4DCT" - pca_components = 7 + + pca_number_of_modes = 6 + # Atlas iterations used to build the reference surface; 1 is a single # template-biased pass. mean_surface_iterations = 3 @@ -126,7 +128,7 @@ workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_surfaces, reference_mesh=reference_surface, - pca_number_of_components=pca_components, + pca_number_of_components=pca_number_of_modes, log_level=log_level, ) @@ -165,7 +167,7 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = min(2, pca_components, len(components), len(eigenvalues)) + mode_count = pca_number_of_modes mode_surface_files: list[Path] = [] xvfb_started = False diff --git a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py index fac26c14..6438afe0 100644 --- a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py @@ -57,6 +57,8 @@ pca_json = tutorial_06_dir / "pca_model.json" pca_mean_file = tutorial_06_dir / "pca_mean_surface.vtp" + pca_number_of_modes = 5 + patient_image_file = repo_root / "data" / "Chest-CT" / "Chest-CT.mha" log_level = logging.INFO @@ -122,10 +124,12 @@ workflow.set_use_pca_registration( use_pca_registration=True, pca_model=pca_model, - pca_number_of_modes=6, + pca_number_of_modes=pca_number_of_modes, use_surface=False, ) + workflow.set_mask_dilation_mm(mask_dilation_mm=40) + # Workflow execution workflow_results = workflow.process() @@ -152,6 +156,14 @@ str(output_dir / f"{project_name}_template_surface_registered.vtp") ) + registered_pca_surface = workflow.pca_template_model_surface + assert registered_pca_surface is not None, ( + "pca_template_model_surface must be set after process()" + ) + registered_pca_surface.save( + str(output_dir / f"{project_name}_pca_surface_registered.vtp") + ) + # Testing TestTools( class_name=project_name, diff --git a/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py b/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py index b5bc566d..24a14252 100644 --- a/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py @@ -15,8 +15,10 @@ filenames and written explicitly into the manifest (the workflow never parses filenames). -2. Split the cases into train / validation / held-out test and train the - MeshGraphNet (``WorkflowTrainPhysicsNeMo`` driving ``TrainPhysicsNeMoMGN``). +2. Split the cases into train and held-out test — plus an optional validation + set, empty by default, which is what makes the intermittent validation RMSE + read ``n/a`` — and train the MeshGraphNet (``WorkflowTrainPhysicsNeMo`` + driving ``TrainPhysicsNeMoMGN``). 3. Evaluate the held-out test cases against their ground-truth phases with :class:`physiotwin4d.WorkflowInferPhysicsNeMo` wrapped in @@ -33,6 +35,15 @@ Edge features (per edge): [rel_x, rel_y, rel_z, distance] (from the mean shape) Output (per vertex): [dx, dy, dz] (displacement in mm) +Runtime +------- +Measured on the full 10-case DIR-Lab set with the Tutorial 6 lung template +(179k points, 1.07M mesh-graph edges): one training step of ``batch_size`` 4 +takes ~430 ms and peaks near 43 GiB of GPU memory, giving ~9 s per epoch and +roughly 4 hours for the 1500 epochs below. Lower ``batch_size``, or call +``training_method.set_num_processor_checkpoint_segments(...)`` to trade compute +for memory, on a smaller card. + Extra Install Required ---------------------- PhysicsNeMo and PyTorch Geometric must be installed:: @@ -111,17 +122,29 @@ def _write_target_mesh( return target_path -def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: +def _write_case_manifest( + case_dir: Path, manifests_dir: Path, logger: logging.Logger +) -> Optional[Path]: """Write a per-case manifest JSON; return its path (or None if incomplete). A case needs a reference SSM surface, a PCA coefficient file, and at least - two respiratory-phase surfaces. + two respiratory-phase surfaces. A case that is missing any of them is skipped + with the reason logged, so a half-finished Tutorial 8 run is distinguishable + from one that never ran. """ case_id = case_dir.name ref_file = case_dir / f"{case_id}_ssm_surface.vtp" pca_file = case_dir / f"{case_id}_ssm_pca_coefficients.json" phase_files = sorted(case_dir.glob(f"{case_id}_T??_ssm_surface.vtp")) - if not ref_file.exists() or not pca_file.exists() or len(phase_files) < 2: + missing = [] + if not ref_file.exists(): + missing.append(f"reference surface {ref_file.name}") + if not pca_file.exists(): + missing.append(f"PCA coefficients {pca_file.name}") + if len(phase_files) < 2: + missing.append(f"at least 2 phase surfaces (found {len(phase_files)})") + if missing: + logger.warning("Skipping %s: missing %s", case_id, "; ".join(missing)) return None manifests_dir.mkdir(parents=True, exist_ok=True) @@ -178,7 +201,9 @@ def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: num_layers = 2 # MLP layers inside each encoder / processor / decoder block # Explicit held-out splits; every other discovered case is used for training. - # Case1 is also the case held out of the Tutorial 2 ICON finetuning. + # Case1 is also the case held out of the Tutorial 2 ICON finetuning. Adding a + # case to val_cases spends it on the intermittent validation RMSE instead of + # training; empty means that RMSE is reported as "n/a". test_cases = ["Case1Pack"] val_cases: list[str] = [] log_level = logging.INFO @@ -198,17 +223,20 @@ def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: "Run tutorials/tutorial_06_lung_create_statistical_model.py first." ) - # Step 1: build one manifest per valid case and partition into splits + # Step 1: build one manifest per valid case and partition into splits. + # DIR-Lab names case 8 "Case8Deploy" while every other case is "Case*Pack", + # so match on "Case*" to avoid silently dropping it. manifests: dict[str, Path] = {} - for case_dir in sorted(data_dir.glob("Case*Pack")): - manifest_path = _write_case_manifest(case_dir, manifests_dir) + for case_dir in sorted(p for p in data_dir.glob("Case*") if p.is_dir()): + manifest_path = _write_case_manifest(case_dir, manifests_dir, logger) if manifest_path is not None: manifests[case_dir.name] = manifest_path if len(manifests) < 3: raise RuntimeError( f"Found only {len(manifests)} valid case(s) under {data_dir}; need at " - "least 3 for a train / val / test split. Run " + "least 3 to hold one out and still train on a population. See the " + "skip reasons logged above, and run " "tutorials/tutorial_08_lung_fit_model_to_4d_patients.py first." ) From b6d90fcb87ff0c42714acd09c384995218bcb4eb Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 12:52:08 -0400 Subject: [PATCH 3/5] ENH: Multiple components 1. Per-structure USD export ENH: Name USD prims per structure and pick materials from those names save_combined_surfaces now tags each cell with a SegmentationLabelIds array so structure identity survives the merge, and ConvertVTKToUSD splits on that array in addition to boundary_labels. ConvertVTKToUSD gains object_names for the static-merge layout; WorkflowConvertVTKToUSD derives them from each mesh's SegmentationLabelNames. anatomy_type now defaults to None, which resolves a material per prim from the prim name, falling back to the object's AnatomyGroup and then to "other" -- so ventricle_left, myocardium, and the great vessels each get their own look instead of one shared heart material. Adds USDAnatomyTools.resolve_anatomy_type so callers can test a name before applying it rather than catching ValueError. 2. Composite transform flattening (bug) BUG: Splice nested composites when chaining registration transforms itk.HDF5TransformIO refuses to write a CompositeTransform holding another CompositeTransform, which every multi-stage registration produced: RegisterImagesGreedy returns an affine+warp composite, and composing a residual onto it nested that composite. _add_transform_flattened splices sub-transforms in at the position their composite occupied, leaving the mapping unchanged since CompositeTransform applies its queue back to front either way. 3. Retire experiment test harness ENH: Drop the experiment test harness; tutorials are the e2e suite Removes tests/test_experiments.py, the "experiment" marker, and --run-experiments. Experiment scripts are exploratory: they assume interactive display, full-resolution parameters, and data layouts that only exist on the author's machine. The test-mode branches in those scripts go away with the harness that drove them, and experiments/ is omitted from coverage. tests/test_tutorials.py --run-tutorials is now the only end-to-end suite; CI comments and READMEs point there. 4. PCA parameter rename (breaking) ENH: Rename PCA count parameters to number_of_pca_components pca_number_of_components, pca_number_of_modes, and --pca-components / --pca-number-of-modes all become number_of_pca_components, so the workflows, the CLIs, and the docs use one name. Breaking change to both the Python and the CLI interface. 5. Distance-map finetuning tutorial ENH: Add tutorial 2 finetuning uniGradICON on lung distance maps RegisterModelsDistanceMaps feeds ICON rasterized signed squared distance maps, not CT intensities, so stock uniGradICON is out of distribution for that stage. The new tutorial finetunes on exactly that representation using DIR-Lab 4D CT lung segmentations, holding Case 1 out for landmark TRE and Dice evaluation. RegisterModelsDistanceMaps.set_icon_weights_path and WorkflowFitStatisticalModelToPatient .set_labelmap_to_labelmap_icon_weights_path plumb the resulting checkpoint into the labelmap-to-labelmap stage; the labelmap-to-image stage keeps stock weights since it registers the image itself. --- .agents/agents/testing.md | 8 +- .github/workflows/README.md | 2 +- .github/workflows/ci.yml | 19 +- .github/workflows/nightly-health.yml | 2 +- AGENTS.md | 6 +- CLAUDE.md | 8 +- docs/cli_scripts/create_statistical_model.rst | 4 +- .../fit_statistical_model_to_patient.rst | 8 +- docs/cli_scripts/vtk_to_usd.rst | 14 +- docs/contributing.rst | 2 +- docs/developer/core.rst | 4 +- docs/testing.rst | 3 +- docs/tutorials.rst | 37 +- .../convert_chop_alterra_valve_to_usd.py | 17 +- .../convert_chop_tpv25_valve_to_usd.py | 15 +- .../1-input_meshes_to_input_surfaces.py | 5 +- .../2-input_surfaces_to_surfaces_aligned.py | 7 +- .../3-registration_based_correspondence.py | 10 +- ...rfaces_aligned_correspond_to_pca_inputs.py | 7 +- .../5-compute_pca_model.py | 8 +- .../Heart-GatedCT_To_USD/1-register_images.py | 56 +- .../2-generate_segmentation.py | 5 +- ...3-transform_dynamic_and_static_contours.py | 4 +- .../simpleware_heart_segmentation.py | 29 +- .../heart_model_to_model_icp_itk.py | 4 +- .../heart_model_to_model_registration_pca.py | 11 +- .../heart_model_to_patient-CHOPValve.py | 9 +- .../heart_model_to_patient.py | 9 +- experiments/README.md | 56 +- .../Reconstruct4DCT/reconstruct_4d_ct.py | 6 +- .../reconstruct_4d_ct_class.py | 127 ++-- pyproject.toml | 3 +- src/physiotwin4d/cli/convert_vtk_to_usd.py | 7 +- .../cli/create_statistical_model.py | 6 +- .../cli/fit_statistical_model_to_patient.py | 8 +- src/physiotwin4d/contour_tools.py | 36 +- src/physiotwin4d/convert_vtk_to_usd.py | 64 +- src/physiotwin4d/register_images_base.py | 39 +- .../register_models_distance_maps.py | 17 + src/physiotwin4d/usd_anatomy_tools.py | 19 + .../workflow_convert_vtk_to_usd.py | 119 +++- .../workflow_create_statistical_model.py | 18 +- ...rkflow_fit_statistical_model_to_patient.py | 42 +- statistics.md | 3 +- tests/README.md | 49 +- tests/conftest.py | 68 +- tests/test_contour_tools.py | 49 ++ tests/test_convert_vtk_to_usd.py | 43 ++ tests/test_experiments.py | 574 ----------------- tests/test_register_images_chain.py | 49 ++ tests/test_tutorials.py | 21 +- tests/test_workflow_convert_vtk_to_usd.py | 173 +++++ tutorials/README.md | 1 + ...orial_02_lung_distancemap_finetune_icon.py | 604 ++++++++++++++++++ tutorials/tutorial_02_lung_finetune_icon.py | 76 ++- tutorials/tutorial_04_heart_ct_to_vtk.py | 11 +- tutorials/tutorial_04_lung_ct_to_vtk.py | 11 +- tutorials/tutorial_05_heart_vtk_to_usd.py | 53 +- ...orial_06_heart_create_statistical_model.py | 8 +- ...torial_06_lung_create_statistical_model.py | 6 +- ...7_lung_fit_statistical_model_to_patient.py | 35 +- ...torial_08_lung_fit_model_to_4d_patients.py | 44 +- 62 files changed, 1655 insertions(+), 1103 deletions(-) delete mode 100644 tests/test_experiments.py create mode 100644 tests/test_workflow_convert_vtk_to_usd.py create mode 100644 tutorials/tutorial_02_lung_distancemap_finetune_icon.py diff --git a/.agents/agents/testing.md b/.agents/agents/testing.md index 2ff0bd39..83353ea5 100644 --- a/.agents/agents/testing.md +++ b/.agents/agents/testing.md @@ -14,7 +14,7 @@ wherever practical. - `tests/baselines/` — stored via Git LFS; fetch with `git lfs pull` - `src/physiotwin4d/test_tools.py` — baseline comparison utilities (`TestTools`) - Markers (all opt-in via `--run-`): `slow`, `requires_gpu`, - `requires_simpleware`, `experiment`, `tutorial`. The `requires_data` marker + `requires_simpleware`, `tutorial`. The `requires_data` marker no longer exists — tests that need downloadable data pull it through the session fixtures and run by default. @@ -29,7 +29,7 @@ python -m pytest tests/ -v # fast, recomm python -m pytest tests/test_contour_tools.py -v # single file python -m pytest tests/test_contour_tools.py::TestContourTools -v # single class python -m pytest tests/ -v --run-slow # opt into slow tests -python -m pytest tests/ -v --run-gpu --run-slow # typical local GPU profile (CI runner adds --run-simpleware --run-experiments --run-tutorials) +python -m pytest tests/ -v --run-gpu --run-slow # typical local GPU profile (CI runner adds --run-simpleware --run-tutorials) python -m pytest tests/ --create-baselines # create missing baselines ``` @@ -59,8 +59,8 @@ python -m pytest tests/ --create-baselines # create missi 8. Prefer storing results in subdirectories under `./results/`. 9. Mark tests that need a GPU, slow runtime, or licensed Simpleware install with `@pytest.mark.requires_gpu`, `@pytest.mark.slow`, or - `@pytest.mark.requires_simpleware`. Mark experiment and tutorial tests - with `@pytest.mark.experiment` or `@pytest.mark.tutorial`. Tests that just + `@pytest.mark.requires_simpleware`. Mark tutorial tests with + `@pytest.mark.tutorial`. Tests that just need downloadable data need no marker. 10. Do not mock segmentation or registration models — test real outputs. 11. No emojis in test files (Windows cp1252 encoding has bitten this project). diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 4e22ee18..ed81f499 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -226,7 +226,7 @@ uv pip install -e ".[test,cuda13,physicsnemo]" pytest tests/ --run-gpu # Enable every --run-* bucket at once (slow, GPU, simpleware, -# physicsnemo, experiments, tutorials) +# physicsnemo, tutorials) pytest tests/ --run-all ``` diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f2316677..ba25d1e5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,7 +13,6 @@ name: CI # - requires_gpu -> --run-gpu # - requires_simpleware -> --run-simpleware (also implies GPU) # - requires_physicsnemo -> --run-physicsnemo (needs the [physicsnemo] extra) -# - experiment -> --run-experiments # - tutorial -> --run-tutorials # --run-all enables every bucket above at once. # Tests that need external data download it automatically via fixtures. @@ -107,13 +106,13 @@ jobs: run: | pip list - - name: Run unit tests (fast, no GPU/slow/experiment) - Ubuntu + - name: Run unit tests (fast, no GPU/slow/tutorial) - Ubuntu if: matrix.os == 'ubuntu-latest' run: | xvfb-run -a --server-args="-screen 0 1024x768x24" \ pytest tests/ -v --cov=physiotwin4d --cov-report=xml --cov-report=term --cov-report=html - - name: Run unit tests (fast, no GPU/slow/experiment) - Windows + - name: Run unit tests (fast, no GPU/slow/tutorial) - Windows if: matrix.os == 'windows-latest' run: | pytest tests/ -v --cov=physiotwin4d --cov-report=xml --cov-report=term --cov-report=html @@ -414,9 +413,9 @@ jobs: # - tests/test_transform_tools.py (depends on slow registration tests) # - tests/test_segment_chest_total_segmentator.py (requires CUDA for TotalSegmentator) # -# Experiment tests (EXTREMELY SLOW - hours to complete): -# - tests/test_experiments.py (runs all notebooks in experiments/ subdirectories) -# These tests are NEVER run in CI/CD and must be run manually +# Tutorial tests (SLOW - hours to complete): +# - tests/test_tutorials.py (runs every script in tutorials/ end-to-end) +# These tests are NEVER run in the PR CI and must be opted into # They execute end-to-end workflows that may take multiple hours # # To run locally: @@ -426,9 +425,9 @@ jobs: # pytest tests/test_register_images_ANTS.py -v --run-slow # # Self-hosted GPU runner enables ALL buckets via --run-all -# (--run-gpu --run-slow --run-simpleware --run-physicsnemo --run-experiments --run-tutorials). +# (--run-gpu --run-slow --run-simpleware --run-physicsnemo --run-tutorials). # That runner installs the [physicsnemo] extra in addition to [test,cuda13]. # -# To run experiment tests (manual only, extremely slow): -# pytest tests/test_experiments.py -v --run-experiments -# pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v --run-experiments +# To run tutorial tests (manual only, slow): +# pytest tests/test_tutorials.py -v --run-tutorials +# pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD -v --run-tutorials diff --git a/.github/workflows/nightly-health.yml b/.github/workflows/nightly-health.yml index 7dd01918..923e69d6 100644 --- a/.github/workflows/nightly-health.yml +++ b/.github/workflows/nightly-health.yml @@ -1,6 +1,6 @@ name: Nightly Health -# Run the full test suite (including experiments) on the GPU runner every night. +# Run the full test suite (including tutorials) on the GPU runner every night. # # Schedule: 07:00 UTC daily ≈ 02:00 EST / 03:00 EDT # diff --git a/AGENTS.md b/AGENTS.md index 00bb4dc0..7f128037 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -81,7 +81,6 @@ python -m pytest tests/ -v --run-slow python -m pytest tests/ -v --run-gpu python -m pytest tests/ -v --run-simpleware python -m pytest tests/ -v --run-physicsnemo -python -m pytest tests/ -v --run-experiments python -m pytest tests/ -v --run-tutorials # Enable every bucket at once (equivalent to passing all --run-* flags) @@ -151,7 +150,7 @@ graphify update . # refresh after code changes (AST-only, no API c `X | None`. - Run `python -m pytest tests/ -v` from the active virtual venv to verify changes. Slow, GPU, Simpleware, - experiment, and tutorial tests are auto-skipped unless their opt-in flag is + and tutorial tests are auto-skipped unless their opt-in flag is passed. - Query the graphify knowledge graph (`graphify query ""`) to locate classes, methods, and signatures before searching manually. @@ -216,8 +215,7 @@ graphify update . # refresh after code changes (AST-only, no API c - Mark tests that need a GPU, a slow runtime, or a licensed Simpleware install with `@pytest.mark.requires_gpu`, `@pytest.mark.slow`, or `@pytest.mark.requires_simpleware`. -- Mark experiment and tutorial tests with `@pytest.mark.experiment` or - `@pytest.mark.tutorial`. +- Mark tutorial tests with `@pytest.mark.tutorial`. - Tests that just need downloadable data need no marker; the fixture chain handles it. - Prefer images from `ROOT/data/test/slicer_heart_small` for tests. diff --git a/CLAUDE.md b/CLAUDE.md index 990efcd5..8ffc5fd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -127,15 +127,17 @@ came from. ## Testing -- Fast tests (recommended for development — slow/GPU/Simpleware/experiment - /tutorial tests are auto-skipped unless their opt-in flag is passed) +- Fast tests (recommended for development — slow/GPU/Simpleware/tutorial + tests are auto-skipped unless their opt-in flag is passed) py -m pytest tests/ -v - Baselines in `tests/baselines/` via Git LFS — run `git lfs pull` after cloning - `tests/conftest.py`: session-scoped fixtures chaining download → convert → segment → register - `src/physiotwin4d/test_tools.py`: baseline comparison utilities (`TestTools`, etc.) - Markers (all opt-in via `--run-`): `slow`, `requires_gpu`, - `requires_simpleware`, `experiment`, `tutorial`. Data-dependent tests no + `requires_simpleware`, `tutorial`. Data-dependent tests no longer use a marker — they pull data through fixtures and run by default. +- `experiments/` scripts are exploratory and are not run as tests; the + `tutorials/` scripts are the optional end-to-end suite - Prefer images from `ROOT/data/test/slicer_heart_small` for tests - Prefer storing results in subdirs `./results/` diff --git a/docs/cli_scripts/create_statistical_model.rst b/docs/cli_scripts/create_statistical_model.rst index 848b46b1..ceb044b3 100644 --- a/docs/cli_scripts/create_statistical_model.rst +++ b/docs/cli_scripts/create_statistical_model.rst @@ -71,7 +71,7 @@ With Custom Parameters --sample-meshes-dir ./meshes \ --reference-mesh average_mesh.vtk \ --output-dir ./pca_output \ - --pca-components 20 + --number-of-pca-components 20 Command-Line Arguments ====================== @@ -94,7 +94,7 @@ Required Arguments Optional Arguments ------------------ -``--pca-components N`` +``--number-of-pca-components N`` Number of PCA components to retain (default: 7). See :class:`physiotwin4d.WorkflowCreateStatisticalModel` for the full API and diff --git a/docs/cli_scripts/fit_statistical_model_to_patient.rst b/docs/cli_scripts/fit_statistical_model_to_patient.rst index 11ff2cfe..5f79ae8a 100644 --- a/docs/cli_scripts/fit_statistical_model_to_patient.rst +++ b/docs/cli_scripts/fit_statistical_model_to_patient.rst @@ -58,7 +58,7 @@ Include statistical shape model fitting: --patient-models lv.vtp rv.vtp myo.vtp \ --patient-image patient_ct.nii.gz \ --pca-json pca_model.json \ - --pca-number-of-modes 10 \ + --number-of-pca-components 10 \ --output-dir ./results Command-Line Arguments @@ -106,8 +106,8 @@ PCA Registration Options ``--pca-json PATH`` Path to PCA JSON file for shape-based registration (optional) -``--pca-number-of-modes NUM`` - Number of PCA modes to use (default: 0, uses all if PCA enabled) +``--number-of-pca-components NUM`` + Number of PCA components to use (default: 0, uses all if PCA enabled) Registration Configuration --------------------------- @@ -162,7 +162,7 @@ Example 2: PCA-Based Registration --patient-models lv.vtp rv.vtp \ --patient-image patient_ct.nii.gz \ --pca-json pca_model.json \ - --pca-number-of-modes 10 \ + --number-of-pca-components 10 \ --output-dir results/pca Output Files diff --git a/docs/cli_scripts/vtk_to_usd.rst b/docs/cli_scripts/vtk_to_usd.rst index 4bdea06f..50f407f8 100644 --- a/docs/cli_scripts/vtk_to_usd.rst +++ b/docs/cli_scripts/vtk_to_usd.rst @@ -36,7 +36,7 @@ Solid color: --appearance solid \ --color 1 0 0 -Anatomy material: +One anatomy material for every mesh: .. code-block:: bash @@ -45,6 +45,18 @@ Anatomy material: --appearance anatomy \ --anatomy-type heart +A material per structure. Omitting ``--anatomy-type`` picks each object's +material from its name, and with ``--static-merge`` the objects are named after +the structures recorded in each file's ``SegmentationLabelNames`` field data — +as written by the image-to-VTK workflow: + +.. code-block:: bash + + physiotwin4d-convert-vtk-to-usd patient_highres_*.vtp \ + --output heart_structures.usd \ + --appearance anatomy \ + --static-merge + Colormap from a VTK point data array: .. code-block:: bash diff --git a/docs/contributing.rst b/docs/contributing.rst index e915b9a6..a20c08bd 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -265,7 +265,7 @@ Run Tests # Run with coverage pytest tests/ --cov=src/physiotwin4d --cov-report=html - # Default invocation auto-skips slow/GPU/Simpleware/experiment/tutorial + # Default invocation auto-skips slow/GPU/Simpleware/tutorial pytest tests/ # Opt into specific buckets diff --git a/docs/developer/core.rst b/docs/developer/core.rst index 3e493b58..8b84f3c6 100644 --- a/docs/developer/core.rst +++ b/docs/developer/core.rst @@ -49,9 +49,9 @@ For most code changes, run: py -m pytest tests/ -v -(Slow / GPU / Simpleware / PhysicsNeMo / experiment / tutorial tests are +(Slow / GPU / Simpleware / PhysicsNeMo / tutorial tests are auto-skipped; opt in with ``--run-slow``, ``--run-gpu``, ``--run-simpleware``, -``--run-physicsnemo``, ``--run-experiments``, ``--run-tutorials``, or use +``--run-physicsnemo``, ``--run-tutorials``, or use ``--run-all`` to enable every bucket at once. Data-dependent tests download their data through the session fixtures and run by default.) diff --git a/docs/testing.rst b/docs/testing.rst index 2b00c34c..33890fe9 100644 --- a/docs/testing.rst +++ b/docs/testing.rst @@ -8,7 +8,7 @@ Run the fast test suite during development: pytest tests/ -v -Slow, GPU, Simpleware, experiment, and tutorial tests are auto-skipped unless +Slow, GPU, Simpleware, and tutorial tests are auto-skipped unless their opt-in flag is passed. Tests that depend on downloadable data fetch it automatically via the session fixtures, so no marker filter is needed for them. @@ -23,7 +23,6 @@ Each ``--run-`` flag enables one marker family: pytest tests/ -v --run-gpu # tests marked 'requires_gpu' pytest tests/ -v --run-simpleware # tests marked 'requires_simpleware' pytest tests/ -v --run-physicsnemo # tests marked 'requires_physicsnemo' - pytest tests/ -v --run-experiments # tests marked 'experiment' pytest tests/ -v --run-tutorials # tests marked 'tutorial' Flags compose. A typical local GPU profile is: diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 5c328e28..13a0b4c4 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -421,9 +421,10 @@ Run python tutorials/tutorial_04_lung_ct_to_vtk.py Outputs - ``patient_surfaces.vtp`` (all anatomy in one mesh), per-group and per-label - ``.vtp`` files, ``patient_labelmap.mha`` and two screenshots, under - ``tutorials/output/tutorial_04_{heart,lung}/``. + ``patient_surfaces.vtp`` (all anatomy in one mesh, with a per-cell + ``SegmentationLabelIds`` array so each cell still names the structure it came + from), per-group and per-label ``.vtp`` files, ``patient_labelmap.mha`` and + two screenshots, under ``tutorials/output/tutorial_04_{heart,lung}/``. Adapt to your data Change the input volume path, then choose the segmenter matching your scan: @@ -444,7 +445,8 @@ Workflow :class:`~physiotwin4d.WorkflowConvertVTKToUSD`. Dataset - Tutorial 4's ``patient_surfaces.vtp`` — no image data, no download. + Tutorial 4's per-structure ``patient_*.vtp`` surfaces — no image data, no + download. Requirements CPU only, seconds to run. The cheapest tutorial in the set. @@ -461,15 +463,22 @@ Inner API usage .. code-block:: python workflow = WorkflowConvertVTKToUSD( - input_meshes=[mesh], + input_meshes=meshes, usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", - anatomy_type="heart", + static_merge=True, separate_by_connectivity=True, ) results = workflow.process() + Each input surface keeps the structure name that + :class:`~physiotwin4d.WorkflowConvertImageToVTK` wrote into its + ``field_data['SegmentationLabelNames']``. That name becomes the USD prim + name and, with ``anatomy_type`` left unset, selects the prim's material — + so the left chambers, right chambers, myocardium and great vessels each get + their own look rather than one shared heart material. + Run .. code-block:: bash @@ -481,11 +490,13 @@ Outputs Adapt to your data ``input_meshes`` takes any list of PyVista meshes — pass one per time point, - in order, for an animated scene instead of a static one, and set - ``frames_per_second`` to control playback. ``appearance="anatomy"`` binds - per-organ materials through :class:`~physiotwin4d.USDAnatomyTools`; use - ``anatomy_type`` to pick the palette. For file-in, file-out conversion - without Python, see :doc:`cli_scripts/vtk_to_usd`. + in order, for an animated scene instead of a static one (drop + ``static_merge``), and set ``frames_per_second`` to control playback. + ``appearance="anatomy"`` binds per-organ materials through + :class:`~physiotwin4d.USDAnatomyTools`; set ``anatomy_type`` to force one + palette onto every object, or ``object_names`` to name the prims yourself. + For file-in, file-out conversion without Python, see + :doc:`cli_scripts/vtk_to_usd`. Tutorial 6: Create a PCA Shape Model ==================================== @@ -533,7 +544,7 @@ Inner API usage workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_surfaces, reference_mesh=reference_surface, - pca_number_of_components=pca_components, + number_of_pca_components=number_of_pca_components, ) result = workflow.process() @@ -553,7 +564,7 @@ Adapt to your data The workflow wants a population of meshes plus one reference; point ``sample_meshes`` at your own cohort and let :class:`~physiotwin4d.WorkflowCreateMeanSurface` build the reference when no - natural template exists. ``pca_number_of_components`` trades fidelity + natural template exists. ``number_of_pca_components`` trades fidelity against cohort size — you need more subjects than modes. The saved ``pca_model.json`` is the portable artifact: Tutorials 7 and 8 and :doc:`cli_scripts/create_statistical_model` all speak it. diff --git a/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py b/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py index f66b87d2..3fa3dd4b 100644 --- a/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py +++ b/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py @@ -30,9 +30,6 @@ from physiotwin4d import ConvertVTKToUSD -# Use as a test -from physiotwin4d.test_tools import TestTools - # Import USDTools for post-processing colormap from physiotwin4d.usd_tools import USDTools @@ -40,11 +37,6 @@ # ## 1. Discover and Organize Time-Series Files # %% -# Set to True to use as a test. Automatically done by -# TestTools.running_as_test() helper function. -test_mode = TestTools.running_as_test() -test_mode_step = 4 - # Define data directories (Alterra only). Anchored to the script's location # so the experiment runs from any working directory. script_dir = Path(__file__).resolve().parent @@ -52,10 +44,7 @@ alterra_dir = data_dir / "Alterra" output_dir = script_dir / "results" / "valve4d-alterra" -if test_mode: - output_usd = output_dir / "alterra_test.usd" -else: - output_usd = output_dir / "alterra_full.usd" +output_usd = output_dir / "alterra_full.usd" colormap_primvar_substrs = ["von_mises_stress"] colormap_name = "jet" # matplotlib colormap name @@ -137,10 +126,6 @@ alterra_files = [file_path for _, file_path in alterra_series] alterra_times = [float(time_step) for time_step, _ in alterra_series] -if test_mode: - alterra_files = alterra_files[::test_mode_step] - alterra_times = alterra_times[::test_mode_step] - print(f"\nConverting to: {output_usd}") print(f"Number of time steps: {len(alterra_times)}") print("\nThis may take several minutes...\n") diff --git a/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py b/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py index ccd8dff0..8b7d33ba 100644 --- a/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py +++ b/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py @@ -29,7 +29,6 @@ from pathlib import Path from physiotwin4d import ConvertVTKToUSD -from physiotwin4d.test_tools import TestTools # Import USDTools for post-processing colormap from physiotwin4d.usd_tools import USDTools @@ -38,11 +37,6 @@ # ## 1. Discover and Organize Time-Series Files # %% -# Set to True to use as a test. Automatically done by -# TestTools.running_as_test() helper function. -quick_run = TestTools.running_as_test() -quick_run_step = 4 - # Define data directories (TPV25 only). Anchored to the script's location # so the experiment runs from any working directory. script_dir = Path(__file__).resolve().parent @@ -50,10 +44,7 @@ tpv25_dir = data_dir / "TPV25" output_dir = script_dir / "results" / "valve4d-tpv25" -if quick_run: - output_usd = output_dir / "tpv25_quick.usd" -else: - output_usd = output_dir / "tpv25_full.usd" +output_usd = output_dir / "tpv25_full.usd" colormap_primvar_substrs = ["von_mises_stress"] colormap_name = "jet" # matplotlib colormap name @@ -134,10 +125,6 @@ tpv25_files = [file_path for _, file_path in tpv25_series] tpv25_times = [float(time_step) for time_step, _ in tpv25_series] -if quick_run: - tpv25_files = tpv25_files[::quick_run_step] - tpv25_times = tpv25_times[::quick_run_step] - print(f"\nConverting to: {output_usd}") print(f"Number of time steps: {len(tpv25_times)}") print("\nThis may take several minutes...\n") diff --git a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py index c98fba56..89b27c77 100644 --- a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py +++ b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py @@ -12,8 +12,6 @@ import pyvista as pv -from physiotwin4d.test_tools import TestTools - _HERE = Path(__file__).parent # %% @@ -76,5 +74,4 @@ plotter = pv.Plotter() plotter.add_mesh(first_surface, color="lightblue", show_edges=True) plotter.add_axes() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() diff --git a/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py b/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py index 0e3d4aec..bf59505d 100644 --- a/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py +++ b/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py @@ -22,7 +22,6 @@ import pyvista as pv from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools from physiotwin4d.register_models_icp import RegisterModelsICP _HERE = Path(__file__).parent @@ -168,8 +167,7 @@ plotter.show_axes() plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## 6. Calculate Registration Statistics @@ -266,7 +264,6 @@ plt.tight_layout() plt.savefig(output_dir / "registration_statistics.png", dpi=150, bbox_inches="tight") -if not TestTools.running_as_test(): - plt.show() +plt.show() print(f"\nPlot saved to: {output_dir / 'registration_statistics.png'}") 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 ffd495ba..034c5c9b 100644 --- a/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py +++ b/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py @@ -27,7 +27,6 @@ from pathlib import Path from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools from physiotwin4d.register_models_distance_maps import RegisterModelsDistanceMaps _HERE = Path(__file__).parent @@ -219,8 +218,7 @@ # Link the camera views so they rotate together plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -274,8 +272,7 @@ plotter.show_axes() plotter.camera_position = "iso" - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% # Save registration statistics @@ -318,8 +315,7 @@ plot_file = output_dir / "registration_statistics.png" plt.savefig(plot_file, dpi=150, bbox_inches="tight") print(f"\nPlot saved to: {plot_file}") - if not TestTools.running_as_test(): - plt.show() + plt.show() else: print("\nNo statistics to plot.") diff --git a/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py b/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py index 26e125e4..d7af7e98 100644 --- a/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py +++ b/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py @@ -6,7 +6,6 @@ import pyvista as pv from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools _HERE = Path(__file__).parent @@ -111,8 +110,7 @@ # Link the camera views so they rotate together plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -166,5 +164,4 @@ plotter.show_axes() plotter.camera_position = "iso" - if not TestTools.running_as_test(): - plotter.show() + plotter.show() diff --git a/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py b/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py index 365309c3..49b7325b 100644 --- a/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py +++ b/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py @@ -21,8 +21,6 @@ from sklearn.decomposition import PCA # SparsePCA, ... -from physiotwin4d.test_tools import TestTools - _HERE = Path(__file__).parent n_components = 15 @@ -242,8 +240,7 @@ def generate_pc_variation(pc_index, std_dev_multiplier=3.0): # Link all three views so camera movements are synchronized plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% # Plot variance explained by each component @@ -266,8 +263,7 @@ def generate_pc_variation(pc_index, std_dev_multiplier=3.0): ax2.grid(True, alpha=0.3) plt.tight_layout() -if not TestTools.running_as_test(): - plt.show() +plt.show() print( f"\nTotal variance captured by {n_components} components: {cumulative_variance[-1] * 100:.2f}%" diff --git a/experiments/Heart-GatedCT_To_USD/1-register_images.py b/experiments/Heart-GatedCT_To_USD/1-register_images.py index 4cf632ef..5cd95f2c 100644 --- a/experiments/Heart-GatedCT_To_USD/1-register_images.py +++ b/experiments/Heart-GatedCT_To_USD/1-register_images.py @@ -7,7 +7,6 @@ from physiotwin4d.segment_chest_total_segmentator_with_contrast import ( SegmentChestTotalSegmentatorWithContrast, ) -from physiotwin4d.test_tools import TestTools from physiotwin4d.transform_tools import TransformTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows @@ -15,14 +14,12 @@ # __name__ == "__main__" guard around the top-level work, that re-import fires # segment() again and Python's spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - test_mode = TestTools.running_as_test() - _HERE = Path(__file__).resolve().parent # %% # Number of cardiac frames and step size; downstream scripts must use the same values. N_FRAMES = 21 - FRAME_STEP = 21 if test_mode else 1 + FRAME_STEP = 1 data_dir = _HERE.parent.parent / "data" / "Slicer-Heart-CT" @@ -34,7 +31,6 @@ # %% seg = SegmentChestTotalSegmentatorWithContrast() - seg.fast_mode = test_mode result = seg.segment(fixed_image) # %% labelmap_mask = result["labelmap"] @@ -79,8 +75,7 @@ # %% reg = RegisterImagesANTS() reg.set_mask_dilation(5) - reg.fast_mode = test_mode - reg.set_number_of_iterations([2, 1, 1] if test_mode else [10, 5, 2]) + reg.set_number_of_iterations([10, 5, 2]) # %% for i in range(0, N_FRAMES, FRAME_STEP): @@ -108,8 +103,6 @@ print(f" Done registering whole image for slice {i:03d}.") inverse_transform = results["inverse_transform"] forward_transform = results["forward_transform"] - whole_image_forward_transform = forward_transform - whole_image_inverse_transform = inverse_transform moving_image_reg = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc @@ -129,28 +122,20 @@ compression=True, ) - # Register the dynamic anatomy mask. In test mode, masked ANTs - # registration (mask applied at every pyramid level) is far more - # expensive than the unmasked whole-image case above; reuse the - # whole-image transform instead of re-registering with a mask. + # Register the dynamic anatomy mask. heart_arr = itk.GetArrayFromImage(heart_mask) contrast_arr = itk.GetArrayFromImage(contrast_mask) major_vessels_arr = itk.GetArrayFromImage(major_vessels_mask) dynamic_anatomy_arr = heart_arr + contrast_arr + major_vessels_arr moving_image_dynamic_anatomy_mask = itk.GetImageFromArray(dynamic_anatomy_arr) moving_image_dynamic_anatomy_mask.CopyInformation(moving_image) - if test_mode: - print(f" Reusing whole-image transform for slice {i:03d} (test mode).") - forward_transform = whole_image_forward_transform - inverse_transform = whole_image_inverse_transform - else: - print(f" Registering dynamic anatomy mask for slice {i:03d}...") - reg.set_fixed_image(fixed_image) - reg.set_fixed_mask(fixed_image_dynamic_anatomy_mask) - results = reg.register(moving_image, moving_image_dynamic_anatomy_mask) - print(f" Done registering dynamic anatomy mask for slice {i:03d}.") - inverse_transform = results["inverse_transform"] - forward_transform = results["forward_transform"] + print(f" Registering dynamic anatomy mask for slice {i:03d}...") + reg.set_fixed_image(fixed_image) + reg.set_fixed_mask(fixed_image_dynamic_anatomy_mask) + results = reg.register(moving_image, moving_image_dynamic_anatomy_mask) + print(f" Done registering dynamic anatomy mask for slice {i:03d}.") + inverse_transform = results["inverse_transform"] + forward_transform = results["forward_transform"] moving_image_reg_dynamic_anatomy = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc @@ -175,7 +160,7 @@ compression=True, ) - # Register the static anatomy mask (same test-mode shortcut as above). + # Register the static anatomy mask. lung_arr = itk.GetArrayFromImage(lung_mask) bone_arr = itk.GetArrayFromImage(bone_mask) other_arr = itk.GetArrayFromImage(other_mask) @@ -183,18 +168,13 @@ lung_arr + bone_arr + other_arr ) moving_image_static_mask.CopyInformation(moving_image) - if test_mode: - print(f" Reusing whole-image transform for slice {i:03d} (test mode).") - forward_transform = whole_image_forward_transform - inverse_transform = whole_image_inverse_transform - else: - print(f" Registering static anatomy mask for slice {i:03d}...") - reg.set_fixed_image(fixed_image) - reg.set_fixed_mask(fixed_image_static_mask) - results = reg.register(moving_image, moving_image_static_mask) - print(f" Done registering static anatomy mask for slice {i:03d}.") - inverse_transform = results["inverse_transform"] - forward_transform = results["forward_transform"] + print(f" Registering static anatomy mask for slice {i:03d}...") + reg.set_fixed_image(fixed_image) + reg.set_fixed_mask(fixed_image_static_mask) + results = reg.register(moving_image, moving_image_static_mask) + print(f" Done registering static anatomy mask for slice {i:03d}.") + inverse_transform = results["inverse_transform"] + forward_transform = results["forward_transform"] moving_image_reg_static = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc diff --git a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py index 74824666..5adcdc67 100644 --- a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py +++ b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py @@ -9,7 +9,6 @@ from physiotwin4d.segment_chest_total_segmentator_with_contrast import ( SegmentChestTotalSegmentatorWithContrast, ) -from physiotwin4d.test_tools import TestTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows # the spawn start method re-imports this script in each child; without the @@ -70,7 +69,6 @@ outname = "slice_max" seg = SegmentChestTotalSegmentatorWithContrast() - seg.fast_mode = TestTools.running_as_test() if re_run_image_segmentation: result = seg.segment(max_image) labelmap_image = result["labelmap"] @@ -162,5 +160,4 @@ ) pl.set_background("black") pl.camera_position = "xy" - if not TestTools.running_as_test(): - pl.show() + pl.show() diff --git a/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py b/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py index 166ddde7..0714bbb8 100644 --- a/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py +++ b/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py @@ -7,7 +7,6 @@ from physiotwin4d import ConvertVTKToUSD from physiotwin4d.contour_tools import ContourTools from physiotwin4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator -from physiotwin4d.test_tools import TestTools from physiotwin4d.usd_anatomy_tools import USDAnatomyTools # Defensive: this script only reads `seg.all_mask_ids` today, but if anyone @@ -15,12 +14,11 @@ # multiprocessing.Pool which re-imports the script on Windows (spawn start # method) and crashes with a spawn-cascade RuntimeError. Guard pre-emptively. if __name__ == "__main__": - test_mode = TestTools.running_as_test() _HERE = os.path.dirname(os.path.abspath(__file__)) # Must match N_FRAMES / FRAME_STEP in 1-register_images.py N_FRAMES = 21 - FRAME_STEP = 21 if test_mode else 1 + FRAME_STEP = 1 # %% output_dir = os.path.join(_HERE, "results") diff --git a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py index 37a66dfe..f1c9ba37 100644 --- a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py +++ b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py @@ -33,7 +33,6 @@ from physiotwin4d.landmark_tools import LandmarkTools from physiotwin4d.segment_heart_simpleware import SegmentHeartSimpleware -from physiotwin4d.test_tools import TestTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -59,18 +58,13 @@ # Load a cardiac CT image for segmentation. This should be a 3D volume containing the heart. # %% -if TestTools.running_as_test(): - input_image_path = os.path.join( - _HERE, "..", "..", "data", "CHOP-Valve4D", "CT", "RVOT28-Dias.nii.gz" - ) -else: - root = tk.Tk() - root.withdraw() - input_image_path = filedialog.askopenfilename( - title="Select a cardiac CT image", - filetypes=[("NIfTI", "*.nii.gz"), ("MetaIO", "*.mhd"), ("All files", "*.*")], - ) - root.destroy() +root = tk.Tk() +root.withdraw() +input_image_path = filedialog.askopenfilename( + title="Select a cardiac CT image", + filetypes=[("NIfTI", "*.nii.gz"), ("MetaIO", "*.mhd"), ("All files", "*.*")], +) +root.destroy() # Load the image try: @@ -119,8 +113,7 @@ axes[2].axis("off") plt.tight_layout() - if not TestTools.running_as_test(): - plt.show() + plt.show() print(f"Image intensity range: [{image_array.min():.1f}, {image_array.max():.1f}]") else: @@ -362,8 +355,7 @@ plt.tight_layout() plt.savefig(os.path.join(output_dir, "segmentation_visualization.png"), dpi=150) - if not TestTools.running_as_test(): - plt.show() + plt.show() print( f"Visualization saved to: {os.path.join(output_dir, 'segmentation_visualization.png')}" @@ -417,8 +409,7 @@ # Save screenshot screenshot_path = os.path.join(output_dir, "3d_visualization.png") - if not TestTools.running_as_test(): - plotter.show(screenshot=screenshot_path) + plotter.show(screenshot=screenshot_path) print(f"3D visualization saved to: {screenshot_path}") else: diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py index 73c29d26..03422961 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py @@ -29,7 +29,6 @@ TransformTools, ) from physiotwin4d.image_tools import ImageTools -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -200,5 +199,4 @@ plotter.add_title("ICP Shape Fitting") plotter.add_axes() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py index 19dc87c7..431b2f17 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py @@ -32,7 +32,6 @@ TransformTools, ) from physiotwin4d.image_tools import ImageTools -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -161,12 +160,10 @@ icp_registrar = RegisterModelsICP(fixed_model=patient_surface) -# Use fewer iterations when run as test (pytest) for faster execution -max_iterations_icp = 100 if TestTools.running_as_test() else 2000 icp_result = icp_registrar.register( transform_type="Affine", moving_model=template_model_surface, - max_iterations=max_iterations_icp, + max_iterations=2000, ) # Get the aligned mesh and transform @@ -325,8 +322,7 @@ plotter.add_axes() plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% [markdown] # ## Visualize PCA Displacement Magnitude @@ -385,8 +381,7 @@ ) plotter.add_title("PCA Signed Displacement on Registered Model") plotter.add_axes() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # Save the mesh with displacement data pca_registered_model_with_displacement.save( 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 cacf30ed..bc4168d4 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 @@ -14,7 +14,6 @@ SegmentHeartSimplewareTrimmedBranches, WorkflowFitStatisticalModelToPatient, ) -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -55,7 +54,7 @@ ) registrar.set_use_pca_registration( - True, pca_model=model_pca_data, pca_number_of_modes=model_pca_n_modes + True, pca_model=model_pca_data, number_of_pca_components=model_pca_n_modes ) registrar.set_use_labelmap_to_labelmap_registration(True) @@ -110,8 +109,7 @@ plotter.add_title("Final Registration") plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -128,8 +126,7 @@ scalar_bar_args={"title": "Deformation (mm)"}, ) plotter.add_title("Deformation Magnitude") - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # Print statistics deformation = registered_model_surface["DeformationMagnitude"] 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 66a846ea..13fa11b2 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 @@ -17,7 +17,6 @@ SegmentChestTotalSegmentator, WorkflowFitStatisticalModelToPatient, ) -from physiotwin4d.test_tools import TestTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows # the spawn start method re-imports this script in each child; without the @@ -137,7 +136,7 @@ patient_image=patient_image, ) registrar.set_use_pca_registration( - True, pca_model=pca_model, pca_number_of_modes=pca_n_modes + True, pca_model=pca_model, number_of_pca_components=pca_n_modes ) registrar.set_use_labelmap_to_image_registration( True, @@ -278,8 +277,7 @@ plotter.add_title("Final Registration") plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -296,8 +294,7 @@ scalar_bar_args={"title": "Deformation (mm)"}, ) plotter.add_title("Deformation Magnitude") - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # Print statistics deformation = registered_surface["DeformationMagnitude"] diff --git a/experiments/README.md b/experiments/README.md index 4e347fac..1c896c56 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -192,52 +192,30 @@ These experiments are **starting points for exploration**, not copy-paste soluti The **CLI commands and implementations in `src/physiotwin4d/cli/`** are the production-quality code you should use and extend for real-world digital twin projects. -## Automated Testing +## Testing -A comprehensive test suite is available to validate all experiment scripts: +Experiment scripts are **not** run as tests. They are exploratory research code: +they assume interactive display, full-resolution parameters, and data layouts +that only exist on the author's machine. -```bash -# Run all experiment tests (EXTREMELY SLOW - may take hours) -# NOTE: Requires --run-experiments flag! -pytest tests/test_experiments.py -v --run-experiments - -# Run a specific experiment subdirectory -pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v -s --run-experiments - -# List all scripts that would be run (without executing) -pytest tests/test_experiments.py::test_list_scripts_in_subdir -v -s --run-experiments +The optional end-to-end test suite runs the `tutorials/` scripts instead: -# Validate experiment directory structure -pytest tests/test_experiments.py::test_experiment_structure -v --run-experiments +```bash +pytest tests/test_tutorials.py -v --run-tutorials ``` -**IMPORTANT:** Experiment tests require the `--run-experiments` flag. Without this flag, they are automatically skipped, even if you run `pytest tests/` or target the test file directly. - -### Test Features - -- **Test-mode flag** - When run as tests (pytest with `--run-experiments`), the runner sets `PHYSIOTWIN_RUNNING_AS_TEST=1`. Scripts can read this (e.g. via `physiotwin4d.test_tools.TestTools.running_as_test()`) and use reduced parameters so test runs stay fast. See [tests/EXPERIMENT_TESTS_GUIDE.md](../tests/EXPERIMENT_TESTS_GUIDE.md#running-as-test-physiotwin_running_as_test). -- **One test per subdirectory** - Each experiment subdirectory gets its own test function -- **Alphanumeric ordering** - Scripts execute in alphanumeric order (e.g., `0-`, `1-`, `2-`) -- **Long timeouts** - Each script has up to 1 hour execution time, tests have multi-hour timeouts -- **Detailed output** - Progress reporting, execution summaries, and failure diagnostics -- **Opt-in only** - Requires `--run-experiments` flag; automatically skipped otherwise -- **Protected from CI/CD** - NEVER runs in automated workflows +Run experiment scripts manually, in the order their filename prefixes imply +(`0-`, `1-`, `2-`, ...), from within their own subdirectory: -### Requirements - -These tests require: -- All dependencies installed (see `pyproject.toml`) -- GPU/CUDA support for most experiments -- Large amounts of disk space and memory -- External data downloads (see individual experiment scripts) - -### Important Notes - -**These tests are extremely long-running** - Plan for multiple hours of execution time - -**Not part of CI/CD** - These tests are excluded from all automated workflows +```bash +cd experiments/Heart-GatedCT_To_USD +py 0-download_and_convert_4d_to_3d.py +py 1-register_images.py +``` -**Resource intensive** - Requires GPU, significant memory, and disk space +Each script requires all dependencies installed (see `pyproject.toml`), +GPU/CUDA support in most cases, large amounts of disk space and memory, and +external data downloads. ## Structure diff --git a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py index a01c3e6b..c6c36cb4 100644 --- a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py +++ b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py @@ -4,7 +4,7 @@ import itk import numpy as np -from physiotwin4d import RegisterImagesGreedy, TestTools, TransformTools +from physiotwin4d import RegisterImagesGreedy, TransformTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -16,10 +16,6 @@ if f.endswith(".mha") and f.startswith("slice_") ] -quick_run = TestTools.running_as_test() -if quick_run: - exit(0) - num_files = len(files) files_indx = list(range(num_files)) reference_image_num = 7 diff --git a/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py b/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py index 931ca99c..a5f1ba97 100644 --- a/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py +++ b/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py @@ -24,7 +24,6 @@ RegisterTimeSeriesImages, TransformTools, ) -from physiotwin4d.test_tools import TestTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -69,44 +68,19 @@ def _build_registrar(method_name: str, iterations=None) -> Optional[RegisterImag print(f"Found {len(files)} slice files") # %% -# Configuration: quick run when executed as test (pytest); full run when manual (set quick_run = True for interactive quick test) -quick_run = TestTools.running_as_test() - -# Select files and parameters based on mode -if quick_run: - print("=== QUICK RUN MODE ===") - total_num_files = len(files) - target_num_files = 2 - if total_num_files == 0: - raise FileNotFoundError(f"No slice_*.mha files found in {data_dir}") - target_num_files = min(target_num_files, total_num_files) - file_step = max(1, total_num_files // target_num_files) - files = files[0:total_num_files:file_step] - files_indx = list(range(0, total_num_files, file_step)) - num_files = len(files) - reference_image_num = num_files // 2 - - # Registration parameters - only Greedy for quick run. ICON and - # Greedy_ICON are exercised by dedicated registration tests elsewhere; - # this experiment validates the reconstruction pipeline, not every - # registration backend. - registration_method_names = ["Greedy"] - number_of_iterations_list = [[2, 1, 1]] # For Greedy -else: - print("=== FULL RUN MODE ===") - num_files = len(files) - files_indx = list(range(num_files)) - reference_image_num = 7 - - # Registration parameters - Greedy_ICON is the recommended method - registration_method_names = [ - "Default" - ] # Use default, or ["Greedy", "ICON", "Greedy_ICON"] - number_of_iterations_list = [None] # [ - # [30, 15, 7, 3], - # 20, # For ICON - # [[30, 15, 7, 3], 20], # For Greedy_ICON - # ] +num_files = len(files) +files_indx = list(range(num_files)) +reference_image_num = 7 + +# Registration parameters - Greedy_ICON is the recommended method +registration_method_names = [ + "Default" +] # Use default, or ["Greedy", "ICON", "Greedy_ICON"] +number_of_iterations_list = [None] # [ +# [30, 15, 7, 3], +# 20, # For ICON +# [[30, 15, 7, 3], 20], # For Greedy_ICON +# ] # Common parameters reference_image_file = os.path.join( @@ -254,41 +228,40 @@ def _build_registrar(method_name: str, iterations=None) -> Optional[RegisterImag print(f" Min loss: {np.min(losses):.6f}") print(f" Max loss: {np.max(losses):.6f}") - if not quick_run: - # Generate grid image for visualization - grid_image = tfm_tools.generate_grid_image(fixed_image, 30, 1) - - print(f"Generating {registration_method_name.upper()} grid visualizations...") - for i, img_indx in enumerate(files_indx): - print(f" Generating grid for slice {img_indx:03d}...") - - # Transform grid with inverse transform (FM) - inverse_grid_image = tfm_tools.transform_image( - grid_image, - inverse_transforms[i], - fixed_image, - ) - itk.imwrite( - inverse_grid_image, - os.path.join( - _RESULTS_DIR, - f"slice_fixed_{registration_method_name}_inverse_grid_{img_indx:03d}.mha", - ), - compression=True, - ) - - # Save displacement field as image - inverse_transform_image = tfm_tools.convert_transform_to_displacement_field( - inverse_transforms[i], - fixed_image, - np_component_type=np.float32, - ) - itk.imwrite( - inverse_transform_image, - os.path.join( - _RESULTS_DIR, - f"slice_{registration_method_name}_inverse_{img_indx:03d}_field.mha", - ), - compression=True, - ) - print(f"Grid visualizations saved for {registration_method_name.upper()}") + # Generate grid image for visualization + grid_image = tfm_tools.generate_grid_image(fixed_image, 30, 1) + + print(f"Generating {registration_method_name.upper()} grid visualizations...") + for i, img_indx in enumerate(files_indx): + print(f" Generating grid for slice {img_indx:03d}...") + + # Transform grid with inverse transform (FM) + inverse_grid_image = tfm_tools.transform_image( + grid_image, + inverse_transforms[i], + fixed_image, + ) + itk.imwrite( + inverse_grid_image, + os.path.join( + _RESULTS_DIR, + f"slice_fixed_{registration_method_name}_inverse_grid_{img_indx:03d}.mha", + ), + compression=True, + ) + + # Save displacement field as image + inverse_transform_image = tfm_tools.convert_transform_to_displacement_field( + inverse_transforms[i], + fixed_image, + np_component_type=np.float32, + ) + itk.imwrite( + inverse_transform_image, + os.path.join( + _RESULTS_DIR, + f"slice_{registration_method_name}_inverse_{img_indx:03d}_field.mha", + ), + compression=True, + ) + print(f"Grid visualizations saved for {registration_method_name.upper()}") diff --git a/pyproject.toml b/pyproject.toml index 8e7033ec..ed8dba5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -327,6 +327,7 @@ module = [ # listed by name; add a line here when adding or renaming a tutorial. "tutorial_01_heart_gated_ct_to_usd", "tutorial_01_lung_gated_ct_to_usd", + "tutorial_02_lung_distancemap_finetune_icon", "tutorial_02_lung_finetune_icon", "tutorial_03_heart_reconstruct_highres_4d_ct", "tutorial_03_lung_reconstruct_highres_4d_ct", @@ -384,7 +385,6 @@ markers = [ "requires_gpu: marks tests that require GPU/CUDA support (opt-in via --run-gpu)", "requires_simpleware: marks tests that require a local Synopsys Simpleware Medical installation (opt-in via --run-simpleware)", "requires_physicsnemo: marks tests that require the optional [physicsnemo] extra (opt-in via --run-physicsnemo)", - "experiment: marks tests that run experiment scripts (extremely slow, manual only)", "tutorial: marks tests that run tutorial scripts (data/GPU gated, manual only)", "xdist_group: marks tests to run in the same pytest-xdist worker (prevents parallel conflicts)" ] @@ -395,6 +395,7 @@ source = ["src/physiotwin4d"] omit = [ "*/tests/*", "*/test_*", + "*/experiments/*", "*/__pycache__/*", "*/venv/*", "*/build/*", diff --git a/src/physiotwin4d/cli/convert_vtk_to_usd.py b/src/physiotwin4d/cli/convert_vtk_to_usd.py index a6d201b9..3971adfc 100644 --- a/src/physiotwin4d/cli/convert_vtk_to_usd.py +++ b/src/physiotwin4d/cli/convert_vtk_to_usd.py @@ -115,8 +115,11 @@ def main() -> int: parser.add_argument( "--anatomy-type", choices=ANATOMY_TYPES, - default="heart", - help="Anatomy material when --appearance anatomy (default: heart)", + default=None, + help=( + "Anatomy material applied to every mesh when --appearance anatomy. " + "Omit to pick a material per object from its name (default)." + ), ) parser.add_argument( "--primvar", diff --git a/src/physiotwin4d/cli/create_statistical_model.py b/src/physiotwin4d/cli/create_statistical_model.py index 863b83fd..547c11aa 100644 --- a/src/physiotwin4d/cli/create_statistical_model.py +++ b/src/physiotwin4d/cli/create_statistical_model.py @@ -39,7 +39,7 @@ def main() -> int: --sample-meshes-dir ./meshes \\ --reference-mesh average_mesh.vtk \\ --output-dir ./pca_model \\ - --pca-components 20 + --number-of-pca-components 20 """, ) @@ -75,7 +75,7 @@ def main() -> int: ) parser.add_argument( - "--pca-components", + "--number-of-pca-components", type=int, default=7, help="Number of PCA components to retain (default: 7)", @@ -142,7 +142,7 @@ def main() -> int: workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_meshes, reference_mesh=reference_mesh, - pca_number_of_components=args.pca_components, + number_of_pca_components=args.number_of_pca_components, ) except (ValueError, RuntimeError) as e: print(f"Error initializing workflow: {e}") diff --git a/src/physiotwin4d/cli/fit_statistical_model_to_patient.py b/src/physiotwin4d/cli/fit_statistical_model_to_patient.py index adc16213..63df8ad8 100644 --- a/src/physiotwin4d/cli/fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/cli/fit_statistical_model_to_patient.py @@ -32,7 +32,7 @@ def main() -> int: --patient-models lv.vtp rv.vtp myo.vtp \\ --patient-image patient_ct.nii.gz \\ --pca-json pca_model.json \\ - --pca-number-of-modes 10 \\ + --number-of-pca-components 10 \\ --output-dir ./results # Enable labelmap-to-image refinement (requires template labelmap and label IDs) @@ -110,10 +110,10 @@ def main() -> int: help="Path to PCA JSON file for shape-based registration (optional)", ) parser.add_argument( - "--pca-number-of-modes", + "--number-of-pca-components", type=int, default=0, - help="Number of PCA modes to use (default: 0, uses all if PCA enabled)", + help="Number of PCA components to use (default: 0, uses all if PCA enabled)", ) # Registration configuration @@ -237,7 +237,7 @@ def main() -> int: workflow.set_use_pca_registration( True, pca_model=pca_model, - pca_number_of_modes=args.pca_number_of_modes, + number_of_pca_components=args.number_of_pca_components, ) workflow.set_use_labelmap_to_labelmap_registration( diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 3d7d2c73..8ec8c0c1 100644 --- a/src/physiotwin4d/contour_tools.py +++ b/src/physiotwin4d/contour_tools.py @@ -627,17 +627,27 @@ def save_combined_surfaces( surface's annotation, enabling colour-by-anatomy rendering in Paraview, PyVista, etc. + It also gains a per-cell ``SegmentationLabelIds`` (int32) array, which + carries each cell's originating label ID so structure identity survives + the merge. Downstream, :class:`ConvertVTKToUSD` splits on this array + when given ``mask_ids``, giving one prim (and one anatomy material) per + structure. A surface whose ``field_data['SegmentationLabelIds']`` does + not hold exactly one ID has no per-cell attribution — that is the case + for the per-group surfaces of :class:`WorkflowConvertImageToVTK`, which + are contoured from a merged binary mask — so its cells are tagged ``0``. + Pass the per-label surfaces (``'label_surfaces'``) to get real IDs. + Per-object ``field_data`` is *not* preserved: it is per-object, so a - single merged mesh cannot carry one value per input surface. The keys - set by :meth:`WorkflowConvertImageToVTK._annotate` are therefore lost: + single merged mesh cannot carry one value per input surface. The + remaining keys set by :meth:`WorkflowConvertImageToVTK._annotate` are + therefore lost: - ``AnatomyGroup`` — group name, e.g. ``'heart'``. - ``SegmentationLabelNames`` — structure names within the group. - - ``SegmentationLabelIds`` — corresponding integer label IDs. - ``AnatomyColor`` — RGB float color (survives indirectly as the per-cell ``Color`` array). - Use :meth:`save_surfaces` instead when structure identity must be + Use :meth:`save_surfaces` instead when structure *names* must be recoverable from the saved files. Args: @@ -656,8 +666,20 @@ def save_combined_surfaces( output_dir = os.path.dirname(output_filename) if output_dir: os.makedirs(output_dir, exist_ok=True) - merged = cast( - pv.PolyData, pv.merge(list(surfaces.values()), merge_points=False) - ) + # Shallow copies so tagging does not add an array to the caller's + # surfaces; the point/cell arrays themselves stay shared. + tagged: list[pv.PolyData] = [] + for surface in surfaces.values(): + label_ids = surface.field_data.get("SegmentationLabelIds") + if label_ids is not None and len(label_ids) == 1: + label_id = int(label_ids[0]) + else: + label_id = 0 + tagged_surface = surface.copy(deep=False) + tagged_surface.cell_data["SegmentationLabelIds"] = np.full( + tagged_surface.n_cells, label_id, dtype=np.int32 + ) + tagged.append(tagged_surface) + merged = cast(pv.PolyData, pv.merge(tagged, merge_points=False)) merged.save(output_filename) return output_filename diff --git a/src/physiotwin4d/convert_vtk_to_usd.py b/src/physiotwin4d/convert_vtk_to_usd.py index ff35ba99..b6df3216 100644 --- a/src/physiotwin4d/convert_vtk_to_usd.py +++ b/src/physiotwin4d/convert_vtk_to_usd.py @@ -41,20 +41,6 @@ validate_time_series_topology, ) -_USD_EXTENSIONS = {".usd", ".usda", ".usdc"} - - -def _split_usd_extension(name: str) -> tuple[str, str]: - """Split a trailing USD extension off ``name``. - - Returns ``(name_without_extension, extension)``. ``extension`` is ``".usd"`` - when ``name`` has no recognized USD extension (``.usd``, ``.usda``, ``.usdc``). - """ - suffix = Path(name).suffix - if suffix.lower() in _USD_EXTENSIONS: - return name[: -len(suffix)], suffix - return name, ".usd" - class ConvertVTKToUSD(PhysioTwin4DBase): """ @@ -98,6 +84,7 @@ def __init__( solid_color: tuple[float, float, float] = (0.8, 0.8, 0.8), static_merge: bool = False, time_codes: Optional[list[float]] = None, + object_names: Optional[Sequence[str]] = None, segmenter: Optional[SegmentAnatomyBase] = None, log_level: int | str = logging.INFO, ) -> None: @@ -125,6 +112,12 @@ def __init__( time_codes: Explicit time codes aligned to input_polydata, used when static_merge is False. If None, uses sequential integers [0, 1, 2, ...]. + object_names: Optional prim names aligned to input_polydata, used + when static_merge is True. If None, objects are named + ``{data_basename}_{index}``. Naming objects after the + structure they hold (e.g. "heart_ventricle_left") makes + the stage self-describing and lets downstream material + assignment key off the prim name. segmenter: Optional SegmentAnatomyBase instance used to classify each mask_ids label into an anatomy group (heart / lung / bone / major_vessels / contrast / soft_tissue / other) so labeled @@ -134,12 +127,18 @@ def __init__( log_level: Logging level Raises: - ValueError: If time_codes is not None and its length does not match - input_polydata, or its values are not non-decreasing. + ValueError: If time_codes or object_names is not None and its length + does not match input_polydata, or if time_codes values are not + non-decreasing. """ super().__init__(class_name=self.__class__.__name__, log_level=log_level) - self.data_basename, _ = _split_usd_extension(data_basename) + suffix = Path(data_basename).suffix + self.data_basename = ( + data_basename[: -len(suffix)] + if suffix.lower() in {".usd", ".usda", ".usdc"} + else data_basename + ) self.input_polydata = list(input_polydata) self.mask_ids = mask_ids self.compute_normals = compute_normals @@ -166,8 +165,16 @@ def __init__( "time_codes must be in non-decreasing order; " "got values that decrease between consecutive frames" ) + if object_names is not None and len(object_names) != len(self.input_polydata): + raise ValueError( + f"object_names length ({len(object_names)}) must match " + f"input_polydata length ({len(self.input_polydata)})" + ) self._is_static_merge: bool = static_merge self._time_codes: Optional[list[float]] = time_codes + self.object_names: Optional[list[str]] = ( + list(object_names) if object_names is not None else None + ) # Pre-converted MeshData for each time step; populated by from_files() so # _convert_unified() can reuse the topology-validation work instead of # calling _vtk_to_mesh_data() a second time. @@ -734,7 +741,11 @@ def _convert_static_merge( ) for i, vtk_mesh in enumerate(self.input_polydata): mesh_data = self._vtk_to_mesh_data(vtk_mesh, i) - frame_name = f"{self.data_basename}_{i}" + frame_name = ( + self.object_names[i] + if self.object_names is not None + else f"{self.data_basename}_{i}" + ) if self.separate_by == "none": parts = [(mesh_data, frame_name)] @@ -935,13 +946,20 @@ def _split_by_labels( if isinstance(vtk_mesh, pv.UnstructuredGrid) and self.convert_to_surface: vtk_mesh = vtk_mesh.extract_surface(algorithm="dataset_surface") - # Get boundary labels - if "boundary_labels" not in vtk_mesh.cell_data: - self.logger.warning("No 'boundary_labels' array found - using unified mesh") + # Get per-cell label IDs. 'SegmentationLabelIds' is written by + # ContourTools.save_combined_surfaces when merging per-label surfaces; + # 'boundary_labels' comes from contouring a multi-label labelmap. + if "SegmentationLabelIds" in vtk_mesh.cell_data: + label_array = vtk_mesh.cell_data["SegmentationLabelIds"] + elif "boundary_labels" in vtk_mesh.cell_data: + label_array = vtk_mesh.cell_data["boundary_labels"] + else: + self.logger.warning( + "No 'SegmentationLabelIds' or 'boundary_labels' array found " + "- using unified mesh" + ) return {"default": self._vtk_to_mesh_data(vtk_mesh, time_idx)} - label_array = vtk_mesh.cell_data["boundary_labels"] - # Create submeshes for each label labeled_meshes = {} for label_id, label_name in mask_ids.items(): diff --git a/src/physiotwin4d/register_images_base.py b/src/physiotwin4d/register_images_base.py index ab496adf..fba61d2d 100644 --- a/src/physiotwin4d/register_images_base.py +++ b/src/physiotwin4d/register_images_base.py @@ -485,8 +485,10 @@ def _compose_with_initial( # mapped by the residual, then by the initial transform, to land in the # original moving image. forward_transform = itk.CompositeTransform[itk.D, 3].New() - forward_transform.AddTransform(initial_forward_transform) - forward_transform.AddTransform(cast(itk.Transform, result["forward_transform"])) + self._add_transform_flattened(forward_transform, initial_forward_transform) + self._add_transform_flattened( + forward_transform, cast(itk.Transform, result["forward_transform"]) + ) # The inverse runs the other way -- a moving-grid sample is mapped by the # initial transform's inverse into the pre-warped frame, then by the @@ -496,8 +498,10 @@ def _compose_with_initial( initial_forward_transform, moving_image ) inverse_transform = itk.CompositeTransform[itk.D, 3].New() - inverse_transform.AddTransform(cast(itk.Transform, result["inverse_transform"])) - inverse_transform.AddTransform(initial_inverse) + self._add_transform_flattened( + inverse_transform, cast(itk.Transform, result["inverse_transform"]) + ) + self._add_transform_flattened(inverse_transform, initial_inverse) return { "forward_transform": forward_transform, @@ -505,6 +509,33 @@ def _compose_with_initial( "loss": result["loss"], } + @staticmethod + def _add_transform_flattened( + composite: itk.CompositeTransform, transform: itk.Transform + ) -> None: + """Append a transform to a composite, splicing in nested composites. + + itk.HDF5TransformIO refuses to write a CompositeTransform that holds + another CompositeTransform ("Composite Transform can only be 1st + transform in a file"), which every multi-stage registration would + otherwise produce: RegisterImagesGreedy already returns an affine+warp + composite, and composing a residual onto it would nest that composite. + + Splicing the sub-transforms in at the position their composite occupied + leaves the mapping unchanged, since itk.CompositeTransform applies its + queue back to front either way. + + The down_cast is required: ITK hands back base-typed ``itkTransformD33`` + Python objects from ``GetInverseTransform()`` and ``GetNthTransform()``, + which carry none of CompositeTransform's methods. + """ + transform = itk.down_cast(transform) + if isinstance(transform, itk.CompositeTransform[itk.D, 3]): + for i in range(transform.GetNumberOfTransforms()): + composite.AddTransform(transform.GetNthTransform(i)) + else: + composite.AddTransform(transform) + def _delegate_to( self, other: "RegisterImagesBase", diff --git a/src/physiotwin4d/register_models_distance_maps.py b/src/physiotwin4d/register_models_distance_maps.py index 4e863949..41867d22 100644 --- a/src/physiotwin4d/register_models_distance_maps.py +++ b/src/physiotwin4d/register_models_distance_maps.py @@ -172,6 +172,23 @@ def __init__( self.inverse_transform: Optional[itk.CompositeTransform] = None # Fixed→moving self.registered_model: Optional[pv.PolyData] = None + def set_icon_weights_path(self, weights_path: str) -> None: + """Use a finetuned uniGradICON checkpoint for the deformable stage. + + The distance maps this class registers are not CT intensities, so stock + uniGradICON weights are out of distribution for them. Weights finetuned + on distance maps, e.g. by + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py``, are + supplied here. + + Args: + weights_path: Path to an existing uniGradICON checkpoint. + + Raises: + FileNotFoundError: If weights_path does not exist. + """ + self.registrar_ICON.set_weights_path(weights_path) + def _create_masks_from_models(self) -> None: """Generate distance maps and binary registration masks from moving and fixed models. diff --git a/src/physiotwin4d/usd_anatomy_tools.py b/src/physiotwin4d/usd_anatomy_tools.py index e99e8f99..7bc73fc5 100644 --- a/src/physiotwin4d/usd_anatomy_tools.py +++ b/src/physiotwin4d/usd_anatomy_tools.py @@ -844,6 +844,25 @@ def _resolve_render_params(self, anatomy_type: str) -> Optional[dict[str, Any]]: return self.render_params[key] return None + def resolve_anatomy_type(self, anatomy_type: str) -> Optional[str]: + """Return the material name *anatomy_type* selects, or ``None``. + + Lets callers test a group/organ name before applying it, instead of + catching the :exc:`ValueError` raised by + :meth:`apply_anatomy_material_to_mesh`. Matching is the same as that + method's (see :meth:`_resolve_render_params`), so ``"kidney_left"`` + returns ``"Kidney"``. + + Args: + anatomy_type: A group/organ name or registered render-params key. + + Returns: + The matching material name (e.g. ``"Heart"``), or ``None`` when + nothing matches. + """ + params = self._resolve_render_params(anatomy_type) + return str(params["name"]) if params is not None else None + def get_anatomy_diffuse_color( self, anatomy_type: str ) -> tuple[float, float, float]: diff --git a/src/physiotwin4d/workflow_convert_vtk_to_usd.py b/src/physiotwin4d/workflow_convert_vtk_to_usd.py index c3346ba9..684b0358 100644 --- a/src/physiotwin4d/workflow_convert_vtk_to_usd.py +++ b/src/physiotwin4d/workflow_convert_vtk_to_usd.py @@ -8,13 +8,14 @@ """ import logging +import re from pathlib import Path -from typing import Any, Literal, Optional, Sequence, Union +from typing import Any, Literal, Mapping, Optional, Sequence, Union import pyvista as pv import vtk -from .convert_vtk_to_usd import ConvertVTKToUSD, _split_usd_extension +from .convert_vtk_to_usd import ConvertVTKToUSD from .physiotwin4d_base import PhysioTwin4DBase from .usd_anatomy_tools import USDAnatomyTools from .usd_tools import USDTools @@ -42,7 +43,8 @@ def __init__( time_codes: Optional[list[float]] = None, appearance: AppearanceKind = "solid", solid_color: tuple[float, float, float] = (0.8, 0.8, 0.8), - anatomy_type: str = "heart", + anatomy_type: Optional[str] = None, + object_names: Optional[Sequence[str]] = None, colormap_primvar: Optional[str] = None, colormap_name: str = "viridis", colormap_intensity_range: Optional[tuple[float, float]] = None, @@ -74,8 +76,21 @@ def __init__( static_merge is False. If None, uses sequential integers [0, 1, 2, ...]. appearance: "solid" | "anatomy" | "colormap". solid_color: RGB in [0,1] when appearance == "solid". - anatomy_type: Anatomy material name when appearance == "anatomy" - (e.g. heart, lung, bone, soft_tissue). + anatomy_type: Anatomy material name applied to every mesh when + appearance == "anatomy" (e.g. heart, lung, bone, soft_tissue). + None (default) instead resolves a material per mesh prim from + that prim's name, so a stage whose objects are named after the + structures they hold gets per-structure materials (e.g. + ventricle_left vs. myocardium). A name matching no material + falls back to the object's ``field_data['AnatomyGroup']`` + (so "rib_left_3" still reaches the bone material) and then to + the "other" material. + object_names: Prim names aligned to input_meshes, used when + static_merge is True. None (default) derives them from each + mesh's ``field_data['SegmentationLabelNames']`` when that holds + exactly one name (as written by + :class:`WorkflowConvertImageToVTK`), and falls back to + ``{usd_project_name}_{index}`` otherwise. colormap_primvar: Primvar name for coloring when appearance == "colormap" (e.g. vtk_point_stress_c0). If None, a candidate is auto-picked when possible. colormap_name: Matplotlib colormap name when appearance == "colormap". @@ -84,9 +99,13 @@ def __init__( """ super().__init__(class_name=self.__class__.__name__, log_level=log_level) self.input_meshes = list(input_meshes) - self.usd_project_name, self._usd_extension = _split_usd_extension( - usd_project_name - ) + suffix = Path(usd_project_name).suffix + if suffix.lower() in {".usd", ".usda", ".usdc"}: + self.usd_project_name = usd_project_name[: -len(suffix)] + self._usd_extension = suffix + else: + self.usd_project_name = usd_project_name + self._usd_extension = ".usd" self.output_directory = Path(output_directory) self.separate_by_connectivity = separate_by_connectivity self.separate_by_cell_type = separate_by_cell_type @@ -97,6 +116,7 @@ def __init__( self.appearance = appearance self.solid_color = solid_color self.anatomy_type = anatomy_type + self.object_names = list(object_names) if object_names is not None else None self.colormap_primvar = colormap_primvar self.colormap_name = colormap_name self.colormap_intensity_range = colormap_intensity_range @@ -106,6 +126,49 @@ def __init__( "separate_by_connectivity and separate_by_cell_type cannot both be True" ) + def _read_object_annotations(self) -> list[tuple[Optional[str], Optional[str]]]: + """Return ``(structure name, anatomy group)`` per input mesh. + + Both come from the annotation :class:`WorkflowConvertImageToVTK` writes + onto each surface: the name from ``field_data['SegmentationLabelNames']`` + when it holds exactly one entry, the group from + ``field_data['AnatomyGroup']``. Either is ``None`` when absent. + """ + annotations: list[tuple[Optional[str], Optional[str]]] = [] + for mesh in self.input_meshes: + if not isinstance(mesh, pv.DataSet): + annotations.append((None, None)) + continue + label_names = mesh.field_data.get("SegmentationLabelNames") + groups = mesh.field_data.get("AnatomyGroup") + name = ( + str(label_names[0]) + if label_names is not None and len(label_names) == 1 + else None + ) + group = str(groups[0]) if groups is not None and len(groups) else None + annotations.append((name, group)) + return annotations + + def _anatomy_candidates( + self, mesh_path: str, object_groups: Mapping[str, str] + ) -> list[str]: + """Return the anatomy names to try for *mesh_path*, best match first. + + With ``anatomy_type`` set, that one name is the only candidate. Without + it, the prim's own name is tried first, then the anatomy group of the + object it came from — so ``"rib_left_3"``, which matches no material of + its own, still lands on the bone material through its group. + """ + if self.anatomy_type is not None: + return [self.anatomy_type] + # Connectivity/cell-type splitting appends "_objectN" to the object + # name; strip it to recover the name object_groups is keyed by. + leaf = mesh_path.rsplit("/", 1)[-1] + object_name = re.sub(r"_object\d+$", "", leaf) + group = object_groups.get(object_name) + return [object_name] if group is None else [object_name, group] + def process(self) -> dict[str, Any]: """ Run the full workflow: convert meshes to USD, then apply the chosen appearance. @@ -147,6 +210,28 @@ def process(self) -> dict[str, Any]: else "none" ) + # Object names only name prims in the static-merge layout; a time + # series writes one prim per part across all frames instead. + annotations = self._read_object_annotations() + object_names = None + if self.static_merge: + object_names = self.object_names + if object_names is None and any(name for name, _ in annotations): + object_names = [ + name or f"{self.usd_project_name}_{index}" + for index, (name, _) in enumerate(annotations) + ] + if object_names is not None: + self.log_info("Naming objects: %s", ", ".join(object_names)) + + # Anatomy group per object name, used as the fallback when the name + # itself matches no material (e.g. "rib_left_3" -> the bone group). + object_groups: dict[str, str] = {} + if object_names is not None: + for object_name, (_, group) in zip(object_names, annotations): + if group is not None: + object_groups[object_name] = group + converter = ConvertVTKToUSD( data_basename=self.usd_project_name, input_polydata=self.input_meshes, @@ -156,6 +241,7 @@ def process(self) -> dict[str, Any]: solid_color=self.solid_color, static_merge=self.static_merge, time_codes=time_codes, + object_names=object_names, log_level=self.log_level, ) stage = converter.convert(str(output_usd)) @@ -191,9 +277,22 @@ def process(self) -> dict[str, Any]: elif self.appearance == "anatomy": anatomy_tools = USDAnatomyTools(stage, log_level=self.log_level) for mesh_path in mesh_paths: - anatomy_tools.apply_anatomy_material_to_mesh( - mesh_path, self.anatomy_type + candidates = self._anatomy_candidates(mesh_path, object_groups) + selected = next( + ( + candidate + for candidate in candidates + if anatomy_tools.resolve_anatomy_type(candidate) is not None + ), + None, ) + if selected is None: + self.log_warning( + "No anatomy material matches %s; using 'other'", + " or ".join(candidates), + ) + selected = "other" + anatomy_tools.apply_anatomy_material_to_mesh(mesh_path, selected) stage.Save() elif self.appearance == "colormap": diff --git a/src/physiotwin4d/workflow_create_statistical_model.py b/src/physiotwin4d/workflow_create_statistical_model.py index 0be75ea9..bf5769d2 100644 --- a/src/physiotwin4d/workflow_create_statistical_model.py +++ b/src/physiotwin4d/workflow_create_statistical_model.py @@ -40,7 +40,7 @@ class WorkflowCreateStatisticalModel(PhysioTwin4DBase): Attributes: sample_meshes (list): List of sample mesh DataSets (.vtk/.vtu/.vtp geometry) reference_mesh (pv.DataSet): Reference mesh; its surface is used for alignment - pca_number_of_components (int): Number of PCA components to retain + number_of_pca_components (int): Number of PCA components to retain reference_spatial_resolution (float): Resolution for reference image from mesh reference_buffer_factor (float): Buffer around mesh for reference image """ @@ -49,7 +49,7 @@ def __init__( self, sample_meshes: list[pv.DataSet], reference_mesh: pv.DataSet, - pca_number_of_components: int = 7, + number_of_pca_components: int = 7, reference_spatial_resolution: float = 1.0, reference_buffer_factor: float = 0.25, solve_for_surface_pca: bool = True, @@ -60,7 +60,7 @@ def __init__( Args: sample_meshes: List of sample mesh DataSets (PyVista PolyData or UnstructuredGrid). reference_mesh: Reference mesh; its surface is used to align all samples. - pca_number_of_components: Number of PCA components. Default 7. + number_of_pca_components: Number of PCA components. Default 7. reference_spatial_resolution: Isotropic resolution (mm) for reference image. Default 1.0. reference_buffer_factor: Buffer factor around mesh for reference image. Default 0.25. solve_for_surface_pca: Whether to reduce the reference mesh to a surface. Default True. @@ -71,7 +71,7 @@ def __init__( ) self.sample_meshes = list(sample_meshes) self.reference_mesh = reference_mesh - self.pca_number_of_components = pca_number_of_components + self.number_of_pca_components = number_of_pca_components self.reference_spatial_resolution = reference_spatial_resolution self.reference_buffer_factor = reference_buffer_factor self.solve_for_surface_pca = solve_for_surface_pca @@ -91,9 +91,9 @@ def __init__( self.pca_mean_surface: Optional[pv.PolyData] = None self.pca_mean_mesh: Optional[pv.DataSet] = None - def set_pca_number_of_components(self, n: int) -> None: + def set_number_of_pca_components(self, n: int) -> None: """Set number of PCA components to retain.""" - self.pca_number_of_components = n + self.number_of_pca_components = n def _step1_extract_surfaces(self) -> None: """Extract reference surface and all sample surfaces (notebook 1).""" @@ -239,11 +239,11 @@ def _step5_compute_pca(self) -> None: raise ValueError( f"At least 2 samples are required for PCA. Got {data_matrix.shape[0]} samples." ) - n_comp = min(self.pca_number_of_components, data_matrix.shape[0] - 1) - if n_comp < self.pca_number_of_components: + n_comp = min(self.number_of_pca_components, data_matrix.shape[0] - 1) + if n_comp < self.number_of_pca_components: self.log_warning( "Reducing PCA components from %d to %d (n_samples=%d)", - self.pca_number_of_components, + self.number_of_pca_components, n_comp, data_matrix.shape[0], ) diff --git a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py index 4b15ec06..e185a76a 100644 --- a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py @@ -84,7 +84,7 @@ class WorkflowFitStatisticalModelToPatient(PhysioTwin4DBase): registrar_Greedy (RegisterImagesGreedy): Greedy registration instance use_pca_registration (bool): Whether PCA registration is enabled (set via set_use_pca_registration) pca_model (dict): PCA model dict when PCA enabled; same structure as WorkflowCreateStatisticalModel output - pca_number_of_modes (int): Number of PCA modes when PCA enabled + number_of_pca_components (int): Number of PCA components when PCA enabled labelmap_interior_object_ids (list): List of labelmap IDs corresponding to interior objects that should not be used when computing a distance map. icp_forward_point_transform : ICP transforms @@ -116,7 +116,7 @@ class WorkflowFitStatisticalModelToPatient(PhysioTwin4DBase): ... ) >>> registrar.set_mask_dilation_mm(20) >>> # To enable PCA registration, call before process(): - >>> # registrar.set_use_pca_registration(True, pca_model=pca_model_dict, pca_number_of_modes=10) + >>> # registrar.set_use_pca_registration(True, pca_model=pca_model_dict, number_of_pca_components=10) >>> # 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.process() @@ -245,6 +245,9 @@ def __init__( # Parameters for labelmap and mask generation self.mask_dilation_mm: float = 10.0 # For binary registration mask generation + # Optional finetuned ICON checkpoint for the labelmap-to-labelmap stage + self.l2l_icon_weights_path: Optional[str] = None + # Stage 1: ICP alignment results self.icp_registrar: Optional[RegisterModelsICP] = None self.icp_inverse_point_transform: Optional[itk.Transform] = None @@ -253,13 +256,13 @@ def __init__( self.icp_template_model_surface: Optional[pv.PolyData] = None self.icp_template_labelmap: Optional[itk.Image] = None - # Stage 1.5: PCA registration results (optional; enable via set_use_pca_registration(True, pca_model, pca_number_of_modes)) + # Stage 1.5: PCA registration results (optional; enable via set_use_pca_registration(True, pca_model, number_of_pca_components)) self.use_pca_registration = False self.pca_registrar: Optional[RegisterModelsPCA] = None self.pca_forward_point_transform: Optional[itk.Transform] = None self.pca_inverse_point_transform: Optional[itk.Transform] = None self.pca_model: Optional[dict[str, Any]] = None - self.pca_number_of_modes: int = 0 + self.number_of_pca_components: int = 0 self.pca_coefficients: Optional[np.ndarray] = None self.pca_template_model: Optional[pv.DataSet] = None self.pca_template_model_surface: Optional[pv.PolyData] = None @@ -296,24 +299,39 @@ def set_mask_dilation_mm(self, mask_dilation_mm: float) -> None: """ self.mask_dilation_mm = mask_dilation_mm + def set_labelmap_to_labelmap_icon_weights_path(self, weights_path: str) -> None: + """Set a finetuned ICON checkpoint for the labelmap-to-labelmap stage. + + That stage (:meth:`register_labelmap_to_labelmap`) registers distance + maps rather than image intensities, so it benefits from weights + finetuned on distance maps -- e.g. by + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py``. The + labelmap-to-image stage keeps the stock weights: it registers the + patient image itself. + + Args: + weights_path: Path to an existing uniGradICON checkpoint. + """ + self.l2l_icon_weights_path = weights_path + def set_use_pca_registration( self, use_pca_registration: bool, pca_model: Optional[dict[str, Any]] = None, - pca_number_of_modes: int = 0, + number_of_pca_components: int = 0, use_surface: bool = False, ) -> None: """Set whether to use PCA-based registration and provide the PCA model. - When enabling (True), pca_model and pca_number_of_modes must be provided. + When enabling (True), pca_model and number_of_pca_components must be provided. Args: use_pca_registration: Whether to use PCA registration after ICP. pca_model: Required when use is True. PCA model dict (e.g. from WorkflowCreateStatisticalModel result["pca_model"]) with keys "eigenvalues" and "components". - pca_number_of_modes: Required when use is True. Number of PCA modes to use. - Default 0 means use all modes. + number_of_pca_components: Required when use is True. Number of PCA + components to use. Default 0 means use all components. use_surface: Whether to use the surface of the patient model for PCA registration. Raises: ValueError: If use is True and pca_model is None. @@ -324,10 +342,10 @@ def set_use_pca_registration( "When enabling PCA registration, pca_model must be provided." ) self.pca_model = pca_model - self.pca_number_of_modes = pca_number_of_modes + self.number_of_pca_components = number_of_pca_components else: self.pca_model = None - self.pca_number_of_modes = 0 + self.number_of_pca_components = 0 self.use_surface = use_surface self.use_pca_registration = use_pca_registration @@ -548,7 +566,7 @@ def register_model_to_model_pca(self) -> dict: self.pca_registrar = RegisterModelsPCA.from_pca_model( pca_template_model=pca_template_model, pca_model=self.pca_model, - pca_number_of_modes=self.pca_number_of_modes, + pca_number_of_modes=self.number_of_pca_components, post_pca_transform=self.icp_forward_point_transform, fixed_model=fixed_model, fixed_distance_map=fixed_distance_map, @@ -688,6 +706,8 @@ def register_labelmap_to_labelmap(self) -> Optional[dict]: mask_dilation_mm=self.mask_dilation_mm, distance_squared_max=(1.25 * self.mask_dilation_mm) ** 2, ) + if self.l2l_icon_weights_path is not None: + labelmap_registrar.set_icon_weights_path(self.l2l_icon_weights_path) # Run deformable registration l2l_result = labelmap_registrar.register( diff --git a/statistics.md b/statistics.md index 82d262ed..40f89c2a 100644 --- a/statistics.md +++ b/statistics.md @@ -152,8 +152,7 @@ PhysioTwin4D operates across several technically demanding domains: - `requires_gpu` - GPU/CUDA-dependent tests (opt-in via `--run-gpu`) - `requires_simpleware` - tests needing a local Synopsys Simpleware Medical install (opt-in via `--run-simpleware`) - `requires_physicsnemo` - tests needing the optional `[physicsnemo]` extra (opt-in via `--run-physicsnemo`) -- `experiment` - runs experiment scripts end-to-end (opt-in via `--run-experiments`; multi-hour) -- `tutorial` - runs tutorial scripts end-to-end (opt-in via `--run-tutorials`) +- `tutorial` - runs tutorial scripts end-to-end (opt-in via `--run-tutorials`; multi-hour) --- diff --git a/tests/README.md b/tests/README.md index 0798ced7..fe7c3a62 100644 --- a/tests/README.md +++ b/tests/README.md @@ -6,9 +6,7 @@ This directory contains comprehensive test suites for the PhysioTwin4D package, - **[TESTING_GUIDE.md](TESTING_GUIDE.md)** - Comprehensive testing guide with setup, troubleshooting, and best practices - **[GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md)** - CI/CD documentation and GitHub Actions workflow details -- **[EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md)** - Guide for running experiment script tests -- **[PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md)** - How parallel execution works with experiment tests -- **[EXPERIMENT_FLAG_USAGE.md](EXPERIMENT_FLAG_USAGE.md)** - Details on the --run-experiments flag +- **[PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md)** - How parallel execution works - **[TEST_FIXES_SUMMARY.md](TEST_FIXES_SUMMARY.md)** - Recent bug fixes and known issues ## Test Categories @@ -33,16 +31,14 @@ This directory contains comprehensive test suites for the PhysioTwin4D package, - **`test_usd_merge.py`** - USD file merging with material preservation - **`test_usd_time_preservation.py`** - Time-varying data validation -### Experiment Tests (EXTREMELY SLOW - Manual Only) -- **`test_experiments.py`** - End-to-end experiment script execution (hours to complete) - - **Opt-in only** - Requires `--run-experiments` flag to run - - **NOT included in CI/CD** - Never runs in automated workflows - - **Automatically skipped** - Won't run with `pytest tests/` unless flag is set - - Runs every `*.py` script in each `experiments/` subdirectory - - Each subdirectory gets its own test - - Scripts run in alphanumeric order - - Requires GPU, CUDA, and all dependencies installed - - See [EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md) for detailed usage instructions +### Tutorial Tests (SLOW - Opt-in) +- **`test_tutorials.py`** - End-to-end execution of each `tutorials/*.py` script, + comparing the screenshots it writes against stored baselines + - **Opt-in only** - Requires `--run-tutorials` flag to run + - Requires GPU, CUDA, all dependencies, and the tutorial datasets + - Scripts in `experiments/` are exploratory and are **not** run as tests + + ## Directory Structure @@ -82,7 +78,7 @@ uv pip install -e ".[test]" ### Run Tests The fast path is the default. Heavy buckets (slow tests, GPU tests, Simpleware -tests, experiment notebooks, tutorial scripts) are **auto-skipped** unless you +tests, tutorial scripts) are **auto-skipped** unless you pass their `--run-` flag. Tests that need downloadable data fetch it through the session fixtures and run by default — there is no `requires_data` marker any more. @@ -105,8 +101,7 @@ Each flag enables one marker family. Flags compose, so you can stack them. | `--run-gpu` | `requires_gpu` | CUDA-dependent tests (ICON, Simpleware, etc.) | | `--run-simpleware` | `requires_simpleware` | Need a licensed Synopsys Simpleware Medical install (also marked `requires_gpu`) | | `--run-physicsnemo` | `requires_physicsnemo` | Need the optional `[physicsnemo]` extra installed | -| `--run-experiments` | `experiment` | End-to-end experiment notebooks (hours to run) | -| `--run-tutorials` | `tutorial` | Tutorial scripts run end-to-end | +| `--run-tutorials` | `tutorial` | Tutorial scripts run end-to-end (hours to run) | | `--run-all` | every bucket above | Equivalent to passing all `--run-*` flags at once | ```bash @@ -122,11 +117,11 @@ pytest tests/ -v --run-all # Full Simpleware coverage (requires Simpleware Medical installed locally) pytest tests/ -v --run-simpleware --run-gpu --run-slow -# Experiment tests (EXTREMELY SLOW — hours to complete) -pytest tests/test_experiments.py -v --run-experiments +# Tutorial tests (SLOW — hours to complete) +pytest tests/test_tutorials.py -v --run-tutorials -# A single experiment by name -pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v -s --run-experiments +# A single tutorial by name +pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD -v -s --run-tutorials ``` ### Common Test Commands @@ -145,7 +140,7 @@ pytest tests/ --create-baselines ## Test Timing Reports -All test runs automatically generate a comprehensive timing report at the end showing individual test durations, session time, and pass/fail/skip counts. The report separates regular tests from experiment tests and highlights the slowest tests. +All test runs automatically generate a comprehensive timing report at the end showing individual test durations, session time, and pass/fail/skip counts. The report separates regular tests from tutorial tests and highlights the slowest tests. ## Test Configuration @@ -163,10 +158,8 @@ All test runs automatically generate a comprehensive timing report at the end sh - `@pytest.mark.requires_physicsnemo` — Tests needing the optional `[physicsnemo]` extra (`pip install "physiotwin4d[physicsnemo]"`, requires Python >= 3.11). Opt in: `--run-physicsnemo`. -- `@pytest.mark.experiment` — End-to-end experiment notebooks (EXTREMELY - SLOW, never in CI). Opt in: `--run-experiments`. -- `@pytest.mark.tutorial` — Tutorial scripts run end-to-end. Opt in: - `--run-tutorials`. +- `@pytest.mark.tutorial` — Tutorial scripts run end-to-end (SLOW, never in + CI). Opt in: `--run-tutorials`. `--run-all` is a convenience flag that turns on every `--run-*` bucket at once. - `@pytest.mark.integration` — Integration tests vs unit tests (filter-only). @@ -201,7 +194,7 @@ Tests automatically run on pull requests via GitHub Actions. The CI workflow: - **Runs fast tests** - USD utilities, data conversion, basic validation - **Skips slow tests** - Registration and segmentation (too slow for CI) -- **Automatically skips experiment tests** - Protected by `--run-experiments` flag requirement +- **Automatically skips tutorial tests** - Protected by `--run-tutorials` flag requirement - **Caches test data** - Speeds up subsequent runs - **Generates coverage** - Reports uploaded to Codecov @@ -209,7 +202,7 @@ Tests automatically run on pull requests via GitHub Actions. The CI workflow: - Platforms: Ubuntu, Windows, macOS - Python versions: 3.11, 3.12 - Target coverage: >70% -- Protection: Experiment tests require `--run-experiments` flag (never used in CI/CD) +- Protection: Tutorial tests require `--run-tutorials` flag (never used in CI/CD) **For detailed CI/CD information**, see [GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md) @@ -269,9 +262,7 @@ Tests automatically: ## Additional Resources - **Detailed Testing Guide**: [TESTING_GUIDE.md](TESTING_GUIDE.md) -- **Experiment Tests Guide**: [EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md) - **Parallel Execution Guide**: [PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md) -- **Experiment Flag Usage**: [EXPERIMENT_FLAG_USAGE.md](EXPERIMENT_FLAG_USAGE.md) - **CI/CD Documentation**: [GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md) - **Recent Fixes**: [TEST_FIXES_SUMMARY.md](TEST_FIXES_SUMMARY.md) - **Main Project**: [../README.md](../README.md) diff --git a/tests/conftest.py b/tests/conftest.py index 0caf99aa..00af88a3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -39,7 +39,6 @@ _RUN_BUCKET_FLAGS = ( - "--run-experiments", "--run-tutorials", "--run-simpleware", "--run-slow", @@ -58,12 +57,6 @@ def _run_bucket_enabled(config: pytest.Config, flag: str) -> bool: def pytest_addoption(parser: pytest.Parser) -> None: """Add custom command-line options for pytest.""" - parser.addoption( - "--run-experiments", - action="store_true", - default=False, - help="Run experiment tests (extremely long-running notebook tests)", - ) parser.addoption( "--run-tutorials", action="store_true", @@ -125,10 +118,6 @@ def pytest_configure(config: pytest.Config) -> None: config.getoption("--create-baselines", default=False) ) - config.addinivalue_line( - "markers", - "experiment: marks tests that run experiment notebooks (extremely slow, manual only)", - ) config.addinivalue_line( "markers", "tutorial: marks tests that run tutorial scripts (data/GPU gated, manual only)", @@ -155,21 +144,12 @@ def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item] ) -> None: """ - Automatically skip experiment and tutorial tests unless their opt-in flags - are passed. + Automatically skip bucketed tests unless their opt-in flags are passed. - This ensures that experiment tests are opt-in only and won't run + This ensures that long-running tests are opt-in only and won't run accidentally when running the normal test suite. """ for item in items: - if "experiment" in item.keywords and not _run_bucket_enabled( - config, "--run-experiments" - ): - item.add_marker( - pytest.mark.skip( - reason="Experiment tests require --run-experiments (or --run-all) to run" - ) - ) if "tutorial" in item.keywords and not _run_bucket_enabled( config, "--run-tutorials" ): @@ -233,7 +213,6 @@ def pytest_runtest_logreport(report: pytest.TestReport) -> None: "nodeid": report.nodeid, "duration": report.duration, "outcome": report.outcome, - "is_experiment": "experiment" in report.keywords, "is_tutorial": "tutorial" in report.keywords, } @@ -250,7 +229,7 @@ def pytest_terminal_summary( Print comprehensive test timing report after all tests complete. This hook is called at the end of the test session to display - timing statistics for all tests, including experiment tests. + timing statistics for all tests, including tutorial tests. """ timings = config._test_timings # type: ignore[attr-defined] tests = timings["tests"] @@ -261,12 +240,9 @@ def pytest_terminal_summary( # Calculate session duration session_duration = datetime.now() - timings["start_time"] - # Separate regular, tutorial, and experiment tests - regular_tests = [ - t for t in tests if not t["is_experiment"] and not t["is_tutorial"] - ] + # Separate regular and tutorial tests + regular_tests = [t for t in tests if not t["is_tutorial"]] tutorial_tests = [t for t in tests if t["is_tutorial"]] - experiment_tests = [t for t in tests if t["is_experiment"]] # Write the timing report terminalreporter.write_sep("=", "TEST TIMING REPORT", bold=True) @@ -323,33 +299,6 @@ def pytest_terminal_summary( ) terminalreporter.write_line("") - # Experiment tests section - if experiment_tests: - terminalreporter.write_sep("-", "Experiment Tests", bold=True) - terminalreporter.write_line(f"Count: {len(experiment_tests)}") - - # Sort by duration (longest first) - sorted_experiments = sorted( - experiment_tests, key=lambda x: x["duration"], reverse=True - ) - - # Calculate total time - experiment_total = sum(t["duration"] for t in experiment_tests) - terminalreporter.write_line( - f"Total Time: {timedelta(seconds=int(experiment_total))}" - ) - terminalreporter.write_line("") - - # Show all experiment tests with timing - terminalreporter.write_line("Individual Test Times:") - for test in sorted_experiments: - outcome_symbol = "+" if test["outcome"] == "passed" else "x" - duration_str = _format_duration(test["duration"]) - terminalreporter.write_line( - f" {outcome_symbol} {duration_str:>10s} {test['nodeid']}" - ) - terminalreporter.write_line("") - # Top 10 slowest tests overall if len(tests) > 10: terminalreporter.write_sep("-", "Top 10 Slowest Tests", bold=True) @@ -358,12 +307,7 @@ def pytest_terminal_summary( for i, test in enumerate(sorted_all, 1): outcome_symbol = "+" if test["outcome"] == "passed" else "x" duration_str = _format_duration(test["duration"]) - if test["is_experiment"]: - test_type = "[EXP]" - elif test["is_tutorial"]: - test_type = "[TUT]" - else: - test_type = "[REG]" + test_type = "[TUT]" if test["is_tutorial"] else "[REG]" terminalreporter.write_line( f" {i:2d}. {outcome_symbol} {duration_str:>10s} {test_type} {test['nodeid']}" ) diff --git a/tests/test_contour_tools.py b/tests/test_contour_tools.py index d01bdbd0..7505aa6b 100644 --- a/tests/test_contour_tools.py +++ b/tests/test_contour_tools.py @@ -337,5 +337,54 @@ def test_contours_from_both_time_points( print(f"Extracted contours from {len(test_labelmaps)} time points") +class TestSaveCombinedSurfaces: + """Structure identity must survive merging into one file.""" + + @staticmethod + def _annotated_sphere(center: tuple[float, float, float], label_id: int) -> Any: + surface = pv.Sphere(radius=1.0, center=center) + surface.field_data["SegmentationLabelIds"] = np.array( + [label_id], dtype=np.int32 + ) + return surface + + def test_per_cell_label_ids_from_single_label_surfaces( + self, tmp_path: Path + ) -> None: + """Each merged cell carries the label ID of the surface it came from.""" + surfaces = { + "atrium_left": self._annotated_sphere((0.0, 0.0, 0.0), 141), + "ventricle_left": self._annotated_sphere((3.0, 0.0, 0.0), 142), + } + + output_file = tmp_path / "combined.vtp" + ContourTools.save_combined_surfaces(surfaces, str(output_file)) + + merged = pv.read(str(output_file)) + label_ids = merged.cell_data["SegmentationLabelIds"] + assert set(np.unique(label_ids)) == {141, 142} + for name, surface in surfaces.items(): + expected = int(surface.field_data["SegmentationLabelIds"][0]) + assert int(np.sum(label_ids == expected)) == surface.n_cells, ( + f"Cell count for {name} did not survive the merge" + ) + # Tagging must not leak back into the caller's surfaces. + for surface in surfaces.values(): + assert "SegmentationLabelIds" not in surface.cell_data + + def test_multi_label_surfaces_tagged_unknown(self, tmp_path: Path) -> None: + """Per-group surfaces list several labels, so no cell can be attributed.""" + group_surface = pv.Sphere(radius=1.0) + group_surface.field_data["SegmentationLabelIds"] = np.array( + [141, 142], dtype=np.int32 + ) + + output_file = tmp_path / "combined_group.vtp" + ContourTools.save_combined_surfaces({"heart": group_surface}, str(output_file)) + + merged = pv.read(str(output_file)) + assert set(np.unique(merged.cell_data["SegmentationLabelIds"])) == {0} + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_convert_vtk_to_usd.py b/tests/test_convert_vtk_to_usd.py index 35c1ec68..2ab2c68d 100644 --- a/tests/test_convert_vtk_to_usd.py +++ b/tests/test_convert_vtk_to_usd.py @@ -477,6 +477,49 @@ def test_static_merge_prim_names_use_data_basename(self, tmp_path: Path) -> None f"{prim_path} should have no time samples but got {samples}" ) + def test_mask_ids_split_on_segmentation_label_ids(self, tmp_path: Path) -> None: + """A merged surface file splits on the array save_combined_surfaces writes.""" + mesh = _make_poly() + mesh.cell_data["SegmentationLabelIds"] = np.array( + [141] * 4 + [142] * 5, dtype=np.int32 + ) + converter = ConvertVTKToUSD( + data_basename="P", + input_polydata=[mesh], + mask_ids={141: "atrium_left", 142: "ventricle_left"}, + ) + stage = converter.convert(str(tmp_path / "out.usd")) + + assert stage.GetPrimAtPath("/World/P/Anatomy/atrium_left").IsValid() + assert stage.GetPrimAtPath("/World/P/Anatomy/ventricle_left").IsValid() + + def test_static_merge_object_names_name_prims(self, tmp_path: Path) -> None: + """object_names replaces the positional {data_basename}_{i} naming.""" + mesh_a, mesh_b = _make_poly(), _make_poly() + converter = ConvertVTKToUSD( + data_basename="Organ", + input_polydata=[mesh_a, mesh_b], + static_merge=True, + object_names=["myocardium", "ventricle_left"], + ) + stage = converter.convert(str(tmp_path / "out.usd")) + + assert stage.GetPrimAtPath("/World/Organ/myocardium").IsValid() + assert stage.GetPrimAtPath("/World/Organ/ventricle_left").IsValid() + assert not stage.GetPrimAtPath("/World/Organ/Organ_0").IsValid(), ( + "Positional naming still present" + ) + + def test_object_names_length_mismatch_raises(self) -> None: + """A short object_names list would silently mis-name prims.""" + with pytest.raises(ValueError, match="object_names length"): + ConvertVTKToUSD( + data_basename="Organ", + input_polydata=[_make_poly(), _make_poly()], + static_merge=True, + object_names=["myocardium"], + ) + # ------------------------------------------------------------------ # Gap D — mask_ids / _convert_with_labels # ------------------------------------------------------------------ diff --git a/tests/test_experiments.py b/tests/test_experiments.py deleted file mode 100644 index 9d73368e..00000000 --- a/tests/test_experiments.py +++ /dev/null @@ -1,574 +0,0 @@ -""" -Test suite for running experiment scripts. - -These tests execute Python scripts in the experiments/ directory. Each subdirectory -in experiments/ gets its own test that runs all scripts in that subdirectory in -alphanumeric order. - -Scripts are Jupytext percent-format files (# %% cell separators), converted from the -original Jupyter notebooks while preserving git history via git mv. - -WARNING: These are EXTREMELY long-running tests that may take hours to complete. -They are NOT part of CI/CD and should only be run manually. - -Usage: - # Run all experiment tests - pytest tests/test_experiments.py -v -m experiment - - # Run a specific experiment subdirectory - pytest tests/test_experiments.py::test_experiment_colormap_vtk_to_usd -v - - # Run with detailed output - pytest tests/test_experiments.py -v -s -m experiment - -Note: These tests require all dependencies installed and GPU/CUDA support for -many of the experiments. -""" - -import os -import subprocess -import sys -from pathlib import Path -from typing import Any - -import pytest - -# Base directories -REPO_ROOT = Path(__file__).parent.parent -EXPERIMENTS_DIR = REPO_ROOT / "experiments" - -# Experiment subdirectories to test (in order of complexity/dependencies) -EXPERIMENT_SUBDIRS = [ - "Colormap-VTK_To_USD", - "Convert_VTK_To_USD", - # 'DisplacementField_To_USD', # Disabled - scripts not ready - "Reconstruct4DCT", - "Heart-VTKSeries_To_USD", - "Heart-GatedCT_To_USD", - "Heart-Create_Statistical_Model", - "Heart-Statistical_Model_To_Patient", - "Lung-GatedCT_To_USD", - # 'Lung-VesselsAirways', # Disabled - scripts not ready -] - - -def get_scripts_in_subdir(subdir_name: str) -> list[Path]: - """ - Get all Python scripts in a subdirectory, sorted alphanumerically. - - Args: - subdir_name: Name of the subdirectory in experiments/ - - Returns: - List of Path objects for .py script files, sorted alphanumerically - """ - subdir = EXPERIMENTS_DIR / subdir_name - if not subdir.exists(): - return [] - - scripts = sorted(subdir.glob("*.py")) - return scripts - - -def execute_script(script_path: Path, timeout: int = 3600) -> dict[str, Any]: - """ - Execute a Python experiment script. - - Args: - script_path: Path to the .py script file - timeout: Maximum execution time in seconds (default: 1 hour) - - Returns: - Dictionary with execution results: - - success: bool - - stdout: str - - stderr: str - - returncode: int - - Raises: - subprocess.TimeoutExpired: If script execution exceeds timeout - """ - print(f"\n{'=' * 80}") - print(f"Executing script: {script_path.name}") - print(f"Path: {script_path}") - print(f"Timeout: {timeout} seconds ({timeout // 60} minutes)") - print(f"{'=' * 80}\n") - - cmd = [sys.executable, str(script_path)] - - # So scripts can use reduced parameters when run as tests - env = os.environ.copy() - env["PHYSIOTWIN_RUNNING_AS_TEST"] = "1" - - try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=timeout * 1.5, # Give extra time for startup overhead - cwd=script_path.parent, # Run in script's directory - env=env, - check=False, - ) - - success = result.returncode == 0 - - if success: - print(f"OK: {script_path.name}") - else: - print(f"FAILED: {script_path.name}") - print(f"Return code: {result.returncode}") - if result.stderr: - print(f"Error output:\n{result.stderr}") - - return { - "success": success, - "stdout": result.stdout, - "stderr": result.stderr, - "returncode": result.returncode, - } - - except subprocess.TimeoutExpired: - print(f"TIMEOUT: {script_path.name}") - print(f"Exceeded: {timeout} seconds") - raise - - -def _heart_statistical_model_pca_prerequisites_met() -> tuple[bool, str]: - """ - Check whether PCA model outputs from Heart-Create_Statistical_Model exist. - - Heart-Statistical_Model_To_Patient scripts expect these artifacts from - the Heart-Create_Statistical_Model experiment (notably 5-compute_pca_model.py). - - Returns: - (True, "") if all required paths exist, else (False, reason_message). - """ - pca_output_dir = ( - EXPERIMENTS_DIR / "Heart-Create_Statistical_Model" / "kcl-heart-model" - ) - pca_json = pca_output_dir / "pca_model.json" - pca_mean_vtp = pca_output_dir / "pca_mean.vtp" - - if not pca_output_dir.is_dir(): - return ( - False, - f"PCA model output directory not found: {pca_output_dir}. " - "Run the Heart-Create_Statistical_Model experiment first " - "(e.g. pytest tests/test_experiments.py::test_experiment_create_statistical_model -v -s --run-experiments).", - ) - if not pca_json.is_file(): - return ( - False, - f"PCA model JSON not found: {pca_json}. " - "Complete Heart-Create_Statistical_Model (including 5-compute_pca_model.py) first.", - ) - if not pca_mean_vtp.is_file(): - return ( - False, - f"PCA mean surface not found: {pca_mean_vtp}. " - "Complete Heart-Create_Statistical_Model (including 5-compute_pca_model.py) first.", - ) - return (True, "") - - -def run_experiment_scripts(subdir_name: str, timeout_per_script: int = 3600) -> None: - """ - Run all Python scripts in an experiment subdirectory in alphanumeric order. - - IMPORTANT: Scripts are executed SEQUENTIALLY in alphanumeric order within - this function. This ensures proper dependency handling even when running - tests with multiple pytest workers (e.g., pytest -n 2). - - The sequential execution is enforced by: - 1. Using a standard Python for loop (not parallelized) - 2. Each script must complete before the next begins - 3. Failures in earlier scripts prevent later ones from running - - Args: - subdir_name: Name of the subdirectory in experiments/ - timeout_per_script: Timeout in seconds for each script (default: 1 hour) - - Raises: - AssertionError: If any script fails to execute successfully - """ - scripts = get_scripts_in_subdir(subdir_name) - - if not scripts: - pytest.skip(f"No scripts found in experiments/{subdir_name}") - - print(f"\n{'#' * 80}") - print(f"# Experiment: {subdir_name}") - print(f"# Found {len(scripts)} script(s)") - print("# Sequential execution enforced (scripts run in order)") - print(f"{'#' * 80}\n") - - failed_scripts: list[dict[str, Any]] = [] - successful_scripts: list[str] = [] - - for i, script in enumerate(scripts, 1): - print(f"\n--- Script {i}/{len(scripts)} ---") - print(f"Sequential execution: script {i} must complete before {i + 1} starts") - - try: - result = execute_script(script, timeout=timeout_per_script) - - if result["success"]: - successful_scripts.append(script.name) - else: - failed_scripts.append( - { - "name": script.name, - "returncode": result["returncode"], - "stderr": result["stderr"], - } - ) - # Stop execution on first failure to maintain dependencies - print(f"\nStopping execution: {script.name} failed") - print("Remaining scripts in this experiment will not run.") - break - - except subprocess.TimeoutExpired: - failed_scripts.append( - { - "name": script.name, - "returncode": -1, - "stderr": f"Timeout after {timeout_per_script} seconds", - } - ) - # Stop execution on timeout - print(f"\nStopping execution: {script.name} timed out") - print("Remaining scripts in this experiment will not run.") - break - - except Exception as e: - failed_scripts.append( - {"name": script.name, "returncode": -2, "stderr": str(e)} - ) - # Stop execution on exception - print(f"\nStopping execution: {script.name} raised exception") - print("Remaining scripts in this experiment will not run.") - break - - # Print summary - print(f"\n{'=' * 80}") - print(f"Experiment Summary: {subdir_name}") - print(f"{'=' * 80}") - print(f"Total scripts: {len(scripts)}") - print(f"Successful: {len(successful_scripts)}") - print(f"Failed: {len(failed_scripts)}") - - if successful_scripts: - print("\nSuccessful scripts:") - for name in successful_scripts: - print(f" - {name}") - - if failed_scripts: - print("\nFailed scripts:") - for failure in failed_scripts: - print(f" - {failure['name']}") - print(f" Return code: {failure['returncode']}") - if failure["stderr"]: - # Print first few lines of error - error_lines = failure["stderr"].split("\n")[:10] - for line in error_lines: - print(f" {line}") - - print(f"{'=' * 80}\n") - - # Assert all scripts succeeded - assert not failed_scripts, ( - f"{len(failed_scripts)} script(s) failed in {subdir_name}: " - f"{[f['name'] for f in failed_scripts]}" - ) - - -# ============================================================================ -# Test Functions - One per Experiment Subdirectory -# ============================================================================ - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group( - name="experiment_colormap" -) # Prevent parallel execution within group -def test_experiment_colormap_vtk_to_usd() -> None: - """ - Test Colormap-VTK_To_USD experiment scripts. - - This experiment demonstrates converting VTK files with colormaps to USD format. - - EXECUTION MODEL: - - Scripts run SEQUENTIALLY in alphanumeric order within this test - - This test function is atomic - pytest-xdist treats it as a single unit - - Multiple experiment tests CAN run in parallel (different subdirectories) - - Scripts within THIS experiment CANNOT run in parallel or out of order - """ - run_experiment_scripts("Colormap-VTK_To_USD", timeout_per_script=3600) - - -# DISABLED - Scripts not ready -# @pytest.mark.experiment -# @pytest.mark.slow -# def test_experiment_displacement_field_to_usd(): -# """ -# Test DisplacementField_To_USD experiment scripts. -# -# This experiment demonstrates converting registration displacement fields to USD -# format for visualization in PhysicsNeMo and Omniverse. -# """ -# run_experiment_scripts('DisplacementField_To_USD', timeout_per_script=3600) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_reconstruct4dct") -def test_experiment_reconstruct_4dct() -> None: - """ - Test Reconstruct4DCT experiment scripts. - - This experiment demonstrates 4D CT reconstruction techniques. - - EXECUTION MODEL: - - Scripts run SEQUENTIALLY in alphanumeric order within this test - - Each script must complete before the next begins - - Failure in one script stops execution of remaining scripts - """ - run_experiment_scripts("Reconstruct4DCT", timeout_per_script=1200) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_heart_vtk") -def test_experiment_heart_vtk_series_to_usd() -> None: - """ - Test Heart-VTKSeries_To_USD experiment scripts. - - This experiment converts heart VTK time series data to USD format. - - EXECUTION ORDER (ENFORCED): - 1. 0-download_and_convert_4d_to_3d.py (downloads data) - 2. 1-heart_vtkseries_to_usd.py (uses downloaded data) - - Each script must complete successfully before the next begins. - """ - run_experiment_scripts("Heart-VTKSeries_To_USD", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_heart_gated_ct") -def test_experiment_heart_gated_ct_to_usd() -> None: - """ - Test Heart-GatedCT_To_USD experiment scripts. - - This is the main cardiac imaging pipeline experiment with strict dependencies. - - EXECUTION ORDER (STRICTLY ENFORCED): - 1. 0-download_and_convert_4d_to_3d.py (downloads and converts data) - 2. 1-register_images.py (registers converted images) - 3. 2-generate_segmentation.py (segments registered images) - 4. 3-transform_dynamic_and_static_contours.py (transforms segmentations) - 5. 4-merge_dynamic_and_static_usd.py (merges into final USD) - - Each script depends on outputs from previous scripts. - Execution stops on first failure to prevent cascading errors. - """ - run_experiment_scripts("Heart-GatedCT_To_USD", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_convert_vtk_to_usd") -def test_experiment_convert_vtk_to_usd() -> None: - """ - Test Convert_VTK_To_USD experiment scripts. - - This experiment demonstrates VTK to USD conversion using the library classes. - - EXECUTION ORDER (ENFORCED): - 1. convert_chop_valve_to_usd.py (converts CHOP valve data) - 2. convert_vtk_to_usd_using_class.py (demonstrates library usage) - - Sequential execution ensures examples build on each other. - """ - run_experiment_scripts("Convert_VTK_To_USD", timeout_per_script=3600) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_create_statistical_model") -def test_experiment_create_statistical_model() -> None: - """ - Test Heart-Create_Statistical_Model experiment scripts. - - This experiment demonstrates creating a PCA statistical shape model from the - KCL Heart Model dataset. - - EXECUTION ORDER (ENFORCED): - 1. 1-input_meshes_to_input_surfaces.py (convert meshes to surfaces) - 2. 2-input_surfaces_to_surfaces_aligned.py (align surfaces) - 3. 3-registration_based_correspondence.py (establish point correspondence) - 4. 4-surfaces_aligned_correspond_to_pca_inputs.py (prepare PCA inputs) - 5. 5-compute_pca_model.py (compute PCA model using sklearn) - - Sequential execution ensures data dependencies are met. - """ - run_experiment_scripts("Heart-Create_Statistical_Model", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_heart_statistical_model") -def test_experiment_heart_statistical_model_to_patient() -> None: - """ - Test Heart-Statistical_Model_To_Patient experiment scripts. - - This experiment demonstrates heart model to patient registration using - statistical shape models (PCA). - - PREREQUISITE: Complete Heart-Create_Statistical_Model experiment first to generate - the PCA model data required for this experiment. - - If PCA outputs (kcl-heart-model/pca_model.json, pca_mean.vtp) are missing, this test - is skipped with a clear message so it can be run in isolation after generating them. - - EXECUTION ORDER (ENFORCED): - 1. heart_model_to_model_icp_itk.py (ICP registration) - 2. heart_model_to_model_registration_pca.py (PCA-based registration) - 3. heart_model_to_patient.py (applies registration to patient) - - Sequential execution ensures registration results are available for subsequent steps. - """ - prereq_met, skip_reason = _heart_statistical_model_pca_prerequisites_met() - if not prereq_met: - pytest.skip(skip_reason) - - run_experiment_scripts( - "Heart-Statistical_Model_To_Patient", timeout_per_script=7200 - ) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_lung_gated_ct") -def test_experiment_lung_gated_ct_to_usd() -> None: - """ - Test Lung-GatedCT_To_USD experiment scripts. - - This is the lung imaging pipeline experiment using DirLab 4DCT data. - - EXECUTION ORDER (STRICTLY ENFORCED): - 1. 0-register_dirlab_4dct.py (registers lung 4DCT data) - 2. 1-make_dirlab_models.py (creates 3D models from registered data) - 3. 2-paint_dirlab_models.py (applies textures/materials to models) - 4. Experiment_ArrangeOnStage.py (arranges models in USD scene) - 5. Experiment_CombineModels.py (combines models into single USD) - 6. Experiment_SegReg.py (segmentation and registration experiments) - 7. Experiment_SubSurfaceScatter.py (applies advanced materials) - - Each script depends on outputs from previous scripts. - Execution is sequential and stops on first failure. - """ - run_experiment_scripts("Lung-GatedCT_To_USD", timeout_per_script=5400) - - -# DISABLED - Scripts not ready -# @pytest.mark.experiment -# @pytest.mark.slow -# @pytest.mark.requires_gpu -# @pytest.mark.requires_data -# def test_experiment_lung_vessels_airways(): -# """ -# Test Lung-VesselsAirways experiment scripts. -# -# This experiment demonstrates specialized vessel and airway segmentation -# using deep learning models. -# Expected scripts (in order): -# - 0-GenData.py -# """ -# run_experiment_scripts('Lung-VesselsAirways', timeout_per_script=3600) - - -# ============================================================================ -# Discovery Test - Validate Experiment Structure -# ============================================================================ - - -@pytest.mark.experiment -def test_experiment_structure() -> None: - """ - Validate the structure of the experiments directory. - - This test checks that: - 1. The experiments directory exists - 2. Each expected subdirectory exists - 3. Each subdirectory contains at least one .py script - """ - assert EXPERIMENTS_DIR.exists(), ( - f"Experiments directory not found: {EXPERIMENTS_DIR}" - ) - - missing_subdirs = [] - empty_subdirs = [] - - for subdir_name in EXPERIMENT_SUBDIRS: - subdir = EXPERIMENTS_DIR / subdir_name - - if not subdir.exists(): - missing_subdirs.append(subdir_name) - continue - - scripts = list(subdir.glob("*.py")) - if not scripts: - empty_subdirs.append(subdir_name) - - # Report findings - if missing_subdirs: - print(f"\nWARNING: Missing subdirectories: {missing_subdirs}") - - if empty_subdirs: - print(f"\nWARNING: Empty subdirectories (no scripts): {empty_subdirs}") - - # Print discovered scripts - print("\nDiscovered Scripts:") - for subdir_name in EXPERIMENT_SUBDIRS: - scripts = get_scripts_in_subdir(subdir_name) - if scripts: - print(f"\n{subdir_name}/ ({len(scripts)} script(s)):") - for s in scripts: - print(f" - {s.name}") - - assert not missing_subdirs, f"Missing subdirectories: {missing_subdirs}" - assert not empty_subdirs, f"Empty subdirectories: {empty_subdirs}" - - -# ============================================================================ -# Helper Test - Script Discovery -# ============================================================================ - - -@pytest.mark.experiment -@pytest.mark.parametrize("subdir_name", EXPERIMENT_SUBDIRS) -def test_list_scripts_in_subdir(subdir_name: str) -> None: - """ - List all scripts in each experiment subdirectory. - - This helper test can be used to preview what scripts will be run - without actually executing them. - - Usage: - pytest tests/test_experiments.py::test_list_scripts_in_subdir -v -s - """ - scripts = get_scripts_in_subdir(subdir_name) - - print(f"\n{subdir_name}/ - {len(scripts)} script(s):") - for i, s in enumerate(scripts, 1): - print(f" {i}. {s.name}") - - assert scripts, f"No scripts found in {subdir_name}" diff --git a/tests/test_register_images_chain.py b/tests/test_register_images_chain.py index 62665dfe..627245e1 100644 --- a/tests/test_register_images_chain.py +++ b/tests/test_register_images_chain.py @@ -98,6 +98,55 @@ def test_chain_refines_previous_stage_result() -> None: assert result["loss"] == 2.0 +class _CompositeRegistrar(_RecordingRegistrar): + """Stub registrar returning a CompositeTransform, as RegisterImagesGreedy + does (its result is an affine plus a displacement field).""" + + def registration_method( + self, + moving_image: itk.Image, + moving_mask: Optional[itk.Image] = None, + moving_labelmap: Optional[itk.Image] = None, + moving_image_pre: Optional[itk.Image] = None, + ) -> dict[str, Union[object, float]]: + """Wrap the parent's translations in single-entry composites.""" + result = super().registration_method( + moving_image, moving_mask, moving_labelmap, moving_image_pre + ) + for key in ("forward_transform", "inverse_transform"): + composite = itk.CompositeTransform[itk.D, 3].New() + composite.AddTransform(cast(itk.Transform, result[key])) + result[key] = composite + return result + + +def test_chain_result_holds_no_nested_composite(tmp_path: Any) -> None: + """A stage returning a CompositeTransform must not end up nested. + + itk.HDF5TransformIO refuses to write a CompositeTransform that holds + another one, so a chain over Greedy (which returns affine+warp composites) + would produce transforms that cannot be saved. + """ + stage1 = _CompositeRegistrar("stage1", 1.0) + stage2 = _CompositeRegistrar("stage2", 2.0) + chain = RegisterImagesChain([stage1, stage2]) + chain.set_fixed_image(_small_image()) + + result = chain.register(_small_image()) + + for key in ("forward_transform", "inverse_transform"): + composed = cast(itk.Transform, result[key]) + assert isinstance(composed, itk.CompositeTransform[itk.D, 3]) + for i in range(composed.GetNumberOfTransforms()): + sub = composed.GetNthTransform(i) + assert "Composite" not in sub.GetNameOfClass() + itk.transformwrite(composed, str(tmp_path / f"{key}.hdf")) + + # Splicing the sub-transforms in must leave the mapping unchanged. + forward = cast(itk.Transform, result["forward_transform"]) + assert list(forward.TransformPoint([0.0, 0.0, 0.0])) == [3.0, 0.0, 0.0] + + def test_chain_propagates_fixed_and_moving_state_to_each_child() -> None: """Every child must see a non-None fixed_image_pre (computed via its own preprocess(), not copied from the chain's no-op preprocess()) and the diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py index d5e2fb64..9f15ca1b 100644 --- a/tests/test_tutorials.py +++ b/tests/test_tutorials.py @@ -202,25 +202,22 @@ class TestTutorial05HeartVTKToUSD: def test_run( self, test_directories: dict[str, Path], test_images: list[Any] ) -> None: - # The script reads this exact path and offers no input override, so - # bootstrap Tutorial 4 rather than pointing it at another surface. - vtk_file = ( - _REPO_ROOT - / "tutorials" - / "output" - / "tutorial_04_heart" - / "patient_surfaces.vtp" - ) - if not vtk_file.exists(): + # The script reads this exact directory and offers no input override, + # so bootstrap Tutorial 4 rather than pointing it at other surfaces. + input_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_04_heart" + if not list(input_dir.glob("patient_*.vtp")): _run_tutorial_script("tutorial_04_heart_ct_to_vtk.py") - assert vtk_file.exists(), ( - f"Tutorial 4 bootstrap did not create the expected surface: {vtk_file}" + assert list(input_dir.glob("patient_*.vtp")), ( + f"Tutorial 4 bootstrap did not create surfaces in: {input_dir}" ) out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05_heart" results = _run_tutorial_script("tutorial_05_heart_vtk_to_usd.py") assert results["usd_file"], "USD file path should not be empty" assert Path(results["usd_file"]).exists(), "USD file should exist" + assert len(results["structures"]) > 1, ( + "Per-structure surfaces expected, so that each becomes its own prim" + ) tt = TestTools( class_name=self._class_name, diff --git a/tests/test_workflow_convert_vtk_to_usd.py b/tests/test_workflow_convert_vtk_to_usd.py new file mode 100644 index 00000000..57da7e3b --- /dev/null +++ b/tests/test_workflow_convert_vtk_to_usd.py @@ -0,0 +1,173 @@ +"""Tests for the appearance and object-naming behavior of the VTK-to-USD workflow. + +Synthetic meshes only - no segmentation or image data required. +""" + +from pathlib import Path + +import numpy as np +import pyvista as pv +import pytest +from pxr import Usd, UsdShade + +from physiotwin4d import WorkflowConvertVTKToUSD + + +def _labeled_sphere( + center: tuple[float, float, float], + label_name: str, + group: str | None = None, +) -> pv.PolyData: + """Return a sphere annotated the way WorkflowConvertImageToVTK annotates one.""" + surface = pv.Sphere(radius=1.0, center=center, theta_resolution=8, phi_resolution=8) + surface.field_data["SegmentationLabelNames"] = np.array([label_name]) + if group is not None: + surface.field_data["AnatomyGroup"] = np.array([group]) + return surface + + +def _bound_material_path(stage: Usd.Stage, mesh_path: str) -> str: + prim = stage.GetPrimAtPath(mesh_path) + assert prim.IsValid(), f"Missing prim: {mesh_path}" + binding = UsdShade.MaterialBindingAPI(prim).GetDirectBinding() + return str(binding.GetMaterialPath()) + + +class TestAnatomyAppearance: + """Per-structure materials must follow the structure names on the meshes.""" + + def test_label_names_drive_prim_names_and_materials(self, tmp_path: Path) -> None: + """Each labeled mesh becomes its own prim with its own anatomy material.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium"), + _labeled_sphere((3.0, 0.0, 0.0), "highres_ventricle_left"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + # separate_by_connectivity defaults to True, and each sphere is a + # single connected component, so every object gains one "_object1" part. + myocardium = _bound_material_path( + stage, "/World/heart/highres_myocardium_object1" + ) + ventricle = _bound_material_path( + stage, "/World/heart/highres_ventricle_left_object1" + ) + assert myocardium.endswith("OmniSurface_Myocardium") + assert ventricle.endswith("OmniSurface_Ventricle_Left") + + def test_explicit_anatomy_type_overrides_names(self, tmp_path: Path) -> None: + """A caller-supplied anatomy_type still paints every object the same.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium"), + _labeled_sphere((3.0, 0.0, 0.0), "highres_ventricle_left"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + anatomy_type="heart", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + for name in ("highres_myocardium", "highres_ventricle_left"): + material = _bound_material_path(stage, f"/World/heart/{name}_object1") + assert material.endswith("OmniSurface_Heart") + + def test_unmatched_name_falls_back_to_group(self, tmp_path: Path) -> None: + """No material is named "rib_left_3", so its anatomy group decides.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "rib_left_3", group="bone"), + _labeled_sphere((3.0, 0.0, 0.0), "vertebrae_T7", group="bone"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="chest", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + for name in ("rib_left_3", "vertebrae_T7"): + material = _bound_material_path(stage, f"/World/chest/{name}_object1") + assert material.endswith("OmniSurface_Bone") + + def test_structure_name_wins_over_group(self, tmp_path: Path) -> None: + """A structure with its own material must not collapse onto its group.""" + mesh = _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium", group="heart") + + workflow = WorkflowConvertVTKToUSD( + input_meshes=[mesh], + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + material = _bound_material_path( + stage, "/World/heart/highres_myocardium_object1" + ) + assert material.endswith("OmniSurface_Myocardium") + + def test_unmatched_name_falls_back_to_other(self, tmp_path: Path) -> None: + """A mesh whose name matches no anatomy still gets a material.""" + mesh = _labeled_sphere((0.0, 0.0, 0.0), "calibration_phantom") + + workflow = WorkflowConvertVTKToUSD( + input_meshes=[mesh], + usd_project_name="scan", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + material = _bound_material_path( + stage, "/World/scan/calibration_phantom_object1" + ) + assert material.endswith("OmniSurface_Other") + + def test_unlabeled_meshes_keep_positional_names(self, tmp_path: Path) -> None: + """Without SegmentationLabelNames, naming stays {project}_{index}.""" + meshes = [ + pv.Sphere(radius=1.0, theta_resolution=8, phi_resolution=8), + pv.Sphere( + radius=1.0, center=(3.0, 0.0, 0.0), theta_resolution=8, phi_resolution=8 + ), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="scan", + output_directory=tmp_path, + appearance="anatomy", + anatomy_type="heart", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + assert stage.GetPrimAtPath("/World/scan/scan_0_object1").IsValid() + assert stage.GetPrimAtPath("/World/scan/scan_1_object1").IsValid() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tutorials/README.md b/tutorials/README.md index ba92cf11..48f9050d 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -28,6 +28,7 @@ current working directory. | 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | | 1 | [tutorial_01_lung_gated_ct_to_usd.py](tutorial_01_lung_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Lung gated 4D CT (prepare first) | | 2 | [tutorial_02_lung_finetune_icon.py](tutorial_02_lung_finetune_icon.py) | `WorkflowFinetuneICONRegistration` | DirLab-4DCT (manual) | +| 2 | [tutorial_02_lung_distancemap_finetune_icon.py](tutorial_02_lung_distancemap_finetune_icon.py) | `WorkflowFinetuneICONRegistration` on lung distance maps | DirLab-4DCT (manual) | | 3 | [tutorial_03_heart_reconstruct_highres_4d_ct.py](tutorial_03_heart_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | Slicer-Heart-CT (prepare first) | | 3 | [tutorial_03_lung_reconstruct_highres_4d_ct.py](tutorial_03_lung_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | | 4 | [tutorial_04_heart_ct_to_vtk.py](tutorial_04_heart_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | diff --git a/tutorials/tutorial_02_lung_distancemap_finetune_icon.py b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py new file mode 100644 index 00000000..6bc22b1e --- /dev/null +++ b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py @@ -0,0 +1,604 @@ +""" +Tutorial 2 (Lung, distance maps): Finetune uniGradICON on lung distance maps + +Purpose +------- +``RegisterModelsDistanceMaps`` -- the labelmap-to-labelmap stage of +``WorkflowFitStatisticalModelToPatient`` -- does not register CT intensities. +It rasterizes a signed squared distance map from each surface, normalizes it to +[-1, 1] and scales it by 1000 so that it fills the CT window uniGradICON +preprocesses with, then hands that pair to ICON. Stock uniGradICON has never +seen such an image, so this tutorial finetunes it on exactly that +representation: distance maps rasterized from the lung surfaces segmented out of +the DIR-Lab 4D CT cases. + +The finetuning cohort is every DIR-Lab case except Case 1, and within each case +only every other respiratory phase (``T00``, ``T20``, ``T40``, ``T60``, ``T80``) +-- half the time points, spanning the full breathing cycle at half the +segmentation cost. Each selected phase is segmented once with +``SegmentNVSegmentCTMRI``; the lung labelmap is kept for uniGradICON's Dice +loss and the lung surfaces are combined and rasterized into the distance map +that serves as the training "image". Segmentation outputs are cached on disk, +so a second run of this tutorial only re-runs the finetuning. + +Accuracy is measured on Case 1, which is never seen during finetuning, by +registering ``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) five +ways: ``RegisterImagesGreedy`` deformable on the distance maps, +``RegisterImagesICON`` on the distance maps with the stock and with the +finetuned weights, ``RegisterImagesGreedyICON`` on the distance maps with the +finetuned weights in its ICON stage -- which separates what the finetuned +network adds on its own from what it adds on top of the Greedy result -- and, +as the intensity-based reference point, stock ``RegisterImagesICON`` on the CT +images themselves. No registration masks are used, so the comparison isolates +the method and the weights. + +The two metrics match ``tutorial_02_lung_finetune_icon.py``. The primary metric +is target registration error: DIR-Lab ships 300 expert landmarks for the extreme +phases (T00 and T50) of every case, so each fixed-image landmark is mapped +through the registration transform and compared, in millimeters, against its +moving-image counterpart. The secondary metric is label overlap: the moving +lung labelmap is warped onto the fixed grid by every transform, so the Dice +scores reflect the transform rather than segmentation variability. The moving +image and labelmap resampled onto the fixed grid without registration supply the +"before registration" reference row for both metrics. + +Reported per method: the mean, standard deviation, 95th percentile and maximum +landmark error in millimeters; the mean, 5th percentile, median, 95th +percentile, minimum and maximum of the per-class Dice scores; the number of +mislabeled voxels; and the wall-clock registration time. + +Finetuning artifacts (dataset JSON, YAML config, checkpoint tree) are written +under ``tutorials/network_weights/icon_dirlab_4dct_distancemap``. The final +checkpoint is ``tutorials/network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch``, the +path returned by ``WorkflowFinetuneICONRegistration.expected_weights_path()`` +and the path ``tutorial_07_lung_fit_statistical_model_to_patient.py`` and +``tutorial_08_lung_fit_model_to_4d_patients.py`` look for. That directory is +deleted at the start of every run, so each run finetunes from scratch; see the +comment above the ``shutil.rmtree`` call for how to reuse a previous run. + +Data Required +------------- +Full data: ``data/DirLab-4DCT`` (all 10 cases, converted to HU ``.mha`` by +``data/DirLab-4DCT/fix_downloaded_data.py``), including the raw +``downloaded_data/Case1Pack/ExtremePhases`` landmark files +Test data: ``data/test/DirLab-4DCT`` +""" + +# Imports +from __future__ import annotations + +import csv +import logging +import shutil +import time +from pathlib import Path +from typing import Any, Optional, cast + +import itk +import numpy as np +import pyvista as pv + +from physiotwin4d import ( + ContourTools, + PhysioTwin4DBase, + RegisterImagesBase, + RegisterImagesGreedy, + RegisterImagesGreedyICON, + RegisterImagesICON, + SegmentNVSegmentCTMRI, + TestTools, + TransformTools, + WorkflowConvertImageToVTK, + WorkflowFinetuneICONRegistration, +) + +# Only run if this script is not imported as a module + +# nnUNetv2 (used inside the segmenter) spawns a multiprocessing.Pool and +# unigradicon finetuning is launched as a subprocess that spawns torch workers; +# on Windows the spawn start method re-imports this script in each child, so all +# top-level work stays under the __name__ == "__main__" guard. +if __name__ == "__main__": + # Data directory specification + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_02_lung_distancemap_finetune_icon" + + output_dir = tutorials_dir / "output" / "tutorial_02_lung_distancemap" + # Segmented labelmaps, lung surfaces, and rasterized distance maps. Cached + # so re-runs skip the segmentation, which dominates this tutorial's runtime. + derived_dir = output_dir / "distance_maps" + baselines_dir = repo_root / "tests" / "baselines" + + # The workflow writes its dataset JSON, YAML config, and checkpoint tree + # under ``weights_dir / finetune_name``. + weights_dir = tutorials_dir / "network_weights" + finetune_name = "icon_dirlab_4dct_distancemap" + + # Distance-map normalization. WorkflowFitStatisticalModelToPatient passes + # (1.25 * mask_dilation_mm) ** 2 as RegisterModelsDistanceMaps' + # distance_squared_max, so the value below fixes the saturation radius of + # every distance map the finetuned weights will ever see. The lung fitting + # tutorials set the same mask_dilation_mm; changing one without the other + # trains on a different image distribution than it infers on. + mask_dilation_mm = 40.0 + distance_squared_max = (1.25 * mask_dilation_mm) ** 2 + + run_finetuning = True + + # Half the respiratory phases per case, taken as every other phase so the + # selection still spans inhale through exhale. + phase_stride = 2 + + test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + number_of_iterations_greedy: Optional[list[int]] = [1, 0] + number_of_iterations_icon = 1 + epochs = 1 + else: + data_dir = repo_root / "data" / "DirLab-4DCT" + number_of_iterations_greedy = [60, 30, 20] + number_of_iterations_icon = 10 + epochs = 200 + + log_level = logging.INFO + reporter = PhysioTwin4DBase(class_name=class_name, log_level=log_level) + + derived_dir.mkdir(parents=True, exist_ok=True) + + # Held-out evaluation pair (Case 1 is excluded from finetuning). T00 and + # T50 are the extreme inhale/exhale phases, the only pair DIR-Lab supplies + # expert landmarks for. + fixed_file = data_dir / "Case1Pack_T50.mha" + moving_file = data_dir / "Case1Pack_T00.mha" + landmark_dir = data_dir / "downloaded_data" / "Case1Pack" / "ExtremePhases" + fixed_landmark_file = landmark_dir / "Case1_300_T50_xyz.txt" + moving_landmark_file = landmark_dir / "Case1_300_T00_xyz.txt" + missing = [ + str(p) + for p in (fixed_file, moving_file, fixed_landmark_file, moving_landmark_file) + if not p.exists() + ] + if missing: + raise FileNotFoundError( + f"Missing DirLab phase images or landmarks: {missing}.\n" + "See data/DirLab-4DCT/README.md for download instructions." + ) + + # Segmentation and distance-map generation + segmentation_workflow = WorkflowConvertImageToVTK( + segmentation_method=SegmentNVSegmentCTMRI(log_level=log_level), + log_level=log_level, + ) + contour_tools = ContourTools(log_level=log_level) + transform_tools = TransformTools() + + def segment_phase(image_file: Path) -> tuple[Path, Path]: + """Segment one phase's lungs and rasterize their distance map. + + Returns the distance map and lung labelmap paths. Both, plus the + combined lung surface, are cached under ``derived_dir``; an existing + pair short-circuits the segmentation. + """ + distance_map_file = derived_dir / f"{image_file.stem}_distance_map.mha" + labelmap_file = derived_dir / f"{image_file.stem}_lung_labelmap.nii.gz" + surface_file = derived_dir / f"{image_file.stem}_lung_surface.vtp" + if distance_map_file.exists() and labelmap_file.exists(): + return distance_map_file, labelmap_file + + reporter.log_info("Segmenting lungs in %s", image_file.name) + image = itk.imread(str(image_file), pixel_type=itk.F) + segmentation_result = segmentation_workflow.process( + input_image=image, + anatomy_groups=["lung"], + extract_label_surfaces=True, + ) + contour_tools.save_combined_surfaces( + segmentation_result["label_surfaces"], str(surface_file) + ) + itk.imwrite( + segmentation_result["labelmap"], str(labelmap_file), compression=True + ) + surface = cast(pv.PolyData, pv.read(str(surface_file))) + # Rasterize the lung surface into ICON's distance-map representation, + # mirroring ``RegisterModelsDistanceMaps._create_masks_from_models`` so + # the finetuning inputs match what that class feeds ICON at inference: + # a signed squared distance normalized to [-1, 1] by + # ``distance_squared_max``, then multiplied by 1000 to fill the + # [-1000, 1000] HU window uniGradICON's CT preprocessing expects. + distance_map = contour_tools.create_distance_map( + surface, + image, + squared_distance=True, + negative_inside=True, + zero_inside=False, + norm_to_max_distance=distance_squared_max, + ) + itk.GetArrayViewFromImage(distance_map)[...] *= 1000 + itk.imwrite(distance_map, str(distance_map_file), compression=True) + return distance_map_file, labelmap_file + + # Finetuning cohort: every case except Case1Pack, every other phase. + # ``Case10Pack_*`` is kept because only the exact ``Case1Pack_`` prefix is + # excluded. + case_phase_files: dict[str, list[Path]] = {} + for path in sorted(data_dir.glob("Case*_T??.mha")): + if path.name.startswith("Case1Pack_"): + continue + case_phase_files.setdefault(path.name.split("_")[0], []).append(path) + if not case_phase_files: + raise FileNotFoundError( + f"No non-Case1 DirLab phase images found under {data_dir}.\n" + "See data/DirLab-4DCT/README.md for download instructions." + ) + case_phase_files = { + case_id: files[::phase_stride] for case_id, files in case_phase_files.items() + } + reporter.log_info( + "Finetuning cohort: %d cases, %d frames (every %d-th phase)", + len(case_phase_files), + sum(len(files) for files in case_phase_files.values()), + phase_stride, + ) + + subject_distance_map_files: list[list[str]] = [] + subject_labelmap_files: list[list[Optional[str]]] = [] + for case_id, phase_files in case_phase_files.items(): + segmented = [segment_phase(path) for path in phase_files] + subject_distance_map_files.append([str(pair[0]) for pair in segmented]) + subject_labelmap_files.append([str(pair[1]) for pair in segmented]) + + # Always finetune from scratch. uniGradICON refuses to overwrite an + # existing experiment directory: it appends "-N" to the name instead + # (``icon_dirlab_4dct_distancemap_model-5``, ...), while + # expected_weights_path() keeps pointing at the original, never-written + # path. Deleting the tree up front keeps the two in agreement. + # + # To reuse a previous run instead, delete the shutil.rmtree call below and + # guard the process() call: + # weights_path = workflow.expected_weights_path() + # if not weights_path.exists(): + # weights_path = workflow.process() + experiment_dir = weights_dir / finetune_name + if run_finetuning: + if experiment_dir.exists(): + reporter.log_info( + "Removing previous finetuning outputs: %s", experiment_dir + ) + shutil.rmtree(experiment_dir) + + # Unlike tutorial_02_lung_finetune_icon.py, the lungs are segmented here + # anyway, so the labelmaps are supplied and uniGradICON's Dice loss stays + # enabled at its default weight. + # + # lncc_sigma matches the sigma RegisterImagesICON uses at inference, so + # finetuning optimizes the similarity this comparison scores. The + # distance maps are already scaled into [-1000, 1000], so the default CT + # window passes them through unclipped. + workflow = WorkflowFinetuneICONRegistration( + subject_image_files=subject_distance_map_files, + output_dir=weights_dir, + finetune_name=finetune_name, + subject_ids=list(case_phase_files.keys()), + subject_labelmap_files=subject_labelmap_files, + epochs=epochs, + lncc_sigma=5, + log_level=log_level, + ) + weights_path = workflow.process() + else: + weights_path = ( + experiment_dir + / f"{finetune_name}_model" + / "checkpoints" + / "network_weights_final.trch" + ) + # Checked here rather than at the first set_weights_path() call, which + # only happens after the greedy and stock-ICON rows have already run. + if not weights_path.exists(): + raise FileNotFoundError( + f"run_finetuning is False but no checkpoint at {weights_path}. " + "Set run_finetuning = True to finetune from scratch." + ) + + # Registration comparison + fixed_image = itk.imread(str(fixed_file), pixel_type=itk.F) + moving_image = itk.imread(str(moving_file), pixel_type=itk.F) + + fixed_distance_map_file, fixed_labelmap_file = segment_phase(fixed_file) + moving_distance_map_file, moving_labelmap_file = segment_phase(moving_file) + fixed_distance_map = itk.imread(str(fixed_distance_map_file), pixel_type=itk.F) + moving_distance_map = itk.imread(str(moving_distance_map_file), pixel_type=itk.F) + fixed_labelmap = itk.imread(str(fixed_labelmap_file)) + moving_labelmap = itk.imread(str(moving_labelmap_file)) + fixed_labels = itk.array_from_image(fixed_labelmap) + + def read_landmarks(landmark_file: Path, image: itk.Image) -> np.ndarray: + """Read a DIR-Lab landmark file as an (N, 3) array of world points. + + Each line holds one 1-based voxel index as ``x y z``. + """ + indices = np.loadtxt(landmark_file, dtype=int) - 1 + return np.array( + [ + image.TransformIndexToPhysicalPoint([int(v) for v in index]) + for index in indices + ] + ) + + fixed_landmarks = read_landmarks(fixed_landmark_file, fixed_image) + moving_landmarks = read_landmarks(moving_landmark_file, moving_image) + + def landmark_metrics(errors_mm: np.ndarray) -> dict[str, Any]: + """Summarize per-landmark target registration errors, in millimeters.""" + return { + "tre_mean": float(errors_mm.mean()), + "tre_std": float(errors_mm.std()), + "tre_p95": float(np.percentile(errors_mm, 95)), + "tre_max": float(errors_mm.max()), + } + + def landmark_errors(transform: itk.Transform) -> np.ndarray: + """Distance from each mapped fixed landmark to its moving counterpart. + + ``forward_transform`` is the resampling transform: it maps points on the + fixed grid back into moving space, which is the direction the landmark + correspondences are defined in. + """ + mapped = np.array( + [transform.TransformPoint(tuple(point)) for point in fixed_landmarks] + ) + return np.asarray(np.linalg.norm(mapped - moving_landmarks, axis=1)) + + def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: + """Per-class Dice summary against the fixed lung labelmap. + + Classes are the union of the two labelmaps' non-zero ids, so a class + found in only one of them scores 0 rather than being dropped. + """ + labels = itk.array_from_image(labelmap) + classes = np.union1d(np.unique(fixed_labels), np.unique(labels)) + classes = classes[classes != 0] + dice = np.array( + [ + 2.0 + * np.count_nonzero((fixed_labels == c) & (labels == c)) + / (np.count_nonzero(fixed_labels == c) + np.count_nonzero(labels == c)) + for c in classes + ] + ) + return { + "n_classes": int(dice.size), + "dice_mean": float(dice.mean()), + "dice_p5": float(np.percentile(dice, 5)), + "dice_median": float(np.median(dice)), + "dice_p95": float(np.percentile(dice, 95)), + "dice_min": float(dice.min()), + "dice_max": float(dice.max()), + "mislabeled_voxels": int(np.count_nonzero(fixed_labels != labels)), + } + + # Reference row: the moving distance map and labelmap on the fixed grid, + # unregistered. + unregistered_distance_map = itk.resample_image_filter( + moving_distance_map, + ReferenceImage=fixed_distance_map, + UseReferenceImage=True, + ) + unregistered_labelmap = itk.resample_image_filter( + moving_labelmap, + Interpolator=itk.NearestNeighborInterpolateImageFunction.New(moving_labelmap), + ReferenceImage=fixed_labelmap, + UseReferenceImage=True, + ) + + registered_distance_maps: dict[str, itk.Image] = { + "unregistered": unregistered_distance_map + } + labelmaps: dict[str, itk.Image] = {"unregistered": unregistered_labelmap} + rows: list[dict[str, Any]] = [ + { + "method": "unregistered", + "input": "-", + "weights": "-", + "registration_time_s": None, + "loss": None, + **landmark_metrics( + np.linalg.norm(fixed_landmarks - moving_landmarks, axis=1) + ), + **overlap_metrics(unregistered_labelmap), + } + ] + # (method, registration input, ICON weights). The last row registers the CT + # images themselves, as the intensity-based reference point for what the + # distance-map methods achieve from surfaces alone. + methods: list[tuple[str, str, Optional[Path]]] = [ + ("greedy_dmap", "distance_map", None), + ("icon_stock_dmap", "distance_map", None), + ("icon_finetuned_dmap", "distance_map", weights_path), + ("greedy_icon_finetuned_dmap", "distance_map", weights_path), + ("icon_stock_ct", "ct", None), + ] + for method_name, method_input, method_weights in methods: + if method_input == "distance_map": + method_fixed, method_moving = fixed_distance_map, moving_distance_map + else: + method_fixed, method_moving = fixed_image, moving_image + + registrar: RegisterImagesBase + if method_name == "greedy_icon_finetuned_dmap": + # Both stages are configured exactly as the standalone "greedy_dmap" + # and "icon_finetuned_dmap" rows above, so this row differs from + # "icon_finetuned_dmap" only by the Greedy transform ICON starts + # from. + chain = RegisterImagesGreedyICON(log_level=log_level) + chain.greedy.set_transform_type("Deformable") + chain.greedy.set_metric("CC") + if number_of_iterations_greedy is not None: + chain.greedy.set_number_of_iterations(number_of_iterations_greedy) + chain.icon.set_number_of_iterations(number_of_iterations_icon) + chain.icon.set_mass_preservation(False) + chain.icon.set_weights_path(str(method_weights)) + registrar = chain + elif method_name.startswith("greedy"): + registrar = RegisterImagesGreedy(log_level=log_level) + registrar.set_transform_type("Deformable") + # CC is what RegisterModelsDistanceMaps uses on distance maps. + registrar.set_metric("CC") + if number_of_iterations_greedy is not None: + registrar.set_number_of_iterations(number_of_iterations_greedy) + else: + registrar = RegisterImagesICON(log_level=log_level) + # None, not 0: icon_registration rejects 0 and takes None to mean + # "no test-time finetuning steps", so the comparison reflects what + # each set of weights predicts rather than per-pair optimization. + registrar.set_number_of_iterations(number_of_iterations_icon) + # Mass preservation models CT density; a distance map carries no + # mass, so it is enabled only on the CT reference row. + registrar.set_mass_preservation(method_input == "ct") + if method_weights is not None: + registrar.set_weights_path(str(method_weights)) + registrar.set_modality("ct") + registrar.set_fixed_image(method_fixed) + + start_time = time.perf_counter() + result = registrar.register(method_moving) + elapsed_s = time.perf_counter() - start_time + + registered_distance_maps[method_name] = transform_tools.transform_image( + moving_distance_map, result["forward_transform"], fixed_distance_map + ) + labelmaps[method_name] = transform_tools.transform_image( + moving_labelmap, + result["forward_transform"], + fixed_labelmap, + interpolation_method="nearest", + ) + rows.append( + { + "method": method_name, + "input": method_input, + "weights": str(method_weights) if method_weights else "-", + "registration_time_s": elapsed_s, + "loss": float(result["loss"]), + **landmark_metrics(landmark_errors(result["forward_transform"])), + **overlap_metrics(labelmaps[method_name]), + } + ) + + # Result saving + itk.imwrite( + fixed_distance_map, str(output_dir / "fixed_distance_map.mha"), compression=True + ) + for method_name, distance_map in registered_distance_maps.items(): + itk.imwrite( + distance_map, + str(output_dir / f"registered_distance_map_{method_name}.mha"), + compression=True, + ) + for method_name, labelmap in labelmaps.items(): + itk.imwrite( + labelmap, + str(output_dir / f"labelmap_{method_name}.mha"), + compression=True, + ) + + summary_file = output_dir / "registration_summary.csv" + with summary_file.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + + # Reporting + reporter.log_info( + "Case1Pack_T00 -> Case1Pack_T50, error at %d expert landmarks, mm", + len(fixed_landmarks), + ) + reporter.log_info( + " %-26s %7s %7s %7s %7s %9s", + "method", + "mean", + "std", + "p95", + "max", + "time_s", + ) + for row in rows: + elapsed = row["registration_time_s"] + reporter.log_info( + " %-26s %7.2f %7.2f %7.2f %7.2f %9s", + row["method"], + row["tre_mean"], + row["tre_std"], + row["tre_p95"], + row["tre_max"], + "-" if elapsed is None else f"{float(elapsed):.1f}", + ) + + reporter.log_info("Per-class Dice of the warped moving labelmap against the fixed") + reporter.log_info( + " %-26s %7s %7s %7s %7s %7s %7s %7s %12s", + "method", + "classes", + "mean", + "p5", + "median", + "p95", + "min", + "max", + "mislabeled", + ) + for row in rows: + reporter.log_info( + " %-26s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", + row["method"], + row["n_classes"], + row["dice_mean"], + row["dice_p5"], + row["dice_median"], + row["dice_p95"], + row["dice_min"], + row["dice_max"], + row["mislabeled_voxels"], + ) + reporter.log_info("Wrote summary: %s", summary_file) + + # Testing + tt = TestTools( + class_name=class_name, + results_dir=output_dir, + baselines_dir=baselines_dir, + log_level=log_level, + ) + + screenshots: list[Path] = [ + tt.save_screenshot_image_slice( + fixed_distance_map, + "fixed_distance_map.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + ) + ] + for method_name, distance_map in registered_distance_maps.items(): + screenshots.append( + tt.save_screenshot_image_slice( + distance_map, + f"registered_distance_map_{method_name}.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + ) + ) + + tutorial_results = { + "weights_path": weights_path, + "registration_metrics": rows, + "labelmaps": labelmaps, + "summary_file": summary_file, + "registered_distance_maps": registered_distance_maps, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_02_lung_finetune_icon.py b/tutorials/tutorial_02_lung_finetune_icon.py index 0d04de6c..2574789f 100644 --- a/tutorials/tutorial_02_lung_finetune_icon.py +++ b/tutorials/tutorial_02_lung_finetune_icon.py @@ -4,11 +4,14 @@ Purpose ------- Finetune uniGradICON on every DIR-Lab 4D CT case except Case 1, then register -``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) three ways: +``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) four ways: ``RegisterImagesGreedy`` alone, deformable, with its default iteration -schedule, and ``RegisterImagesICON`` with the stock uniGradICON weights and -with the finetuned weights. Case 1 is never seen during finetuning, so it is a held-out -evaluation pair. +schedule; ``RegisterImagesICON`` with the stock uniGradICON weights and with +the finetuned weights; and ``RegisterImagesGreedyICON`` -- the same Greedy +stage initializing an ICON stage that carries the finetuned weights -- which +separates what the finetuned network adds on its own from what it adds on top +of a classical affine-plus-deformable initialization. Case 1 is never seen +during finetuning, so it is a held-out evaluation pair. Accuracy is measured two ways. The primary metric is target registration error: DIR-Lab ships 300 expert landmarks for the extreme phases (T00 and T50) @@ -59,6 +62,7 @@ PhysioTwin4DBase, RegisterImagesBase, RegisterImagesGreedy, + RegisterImagesGreedyICON, RegisterImagesICON, SegmentNVSegmentCTMRI, TestTools, @@ -79,6 +83,8 @@ class_name = "tutorial_02_lung_finetune_icon" output_dir = tutorials_dir / "output" / "tutorial_02_lung" + # Segmented labelmaps, cached so re-runs skip the segmentation. + labelmaps_dir = output_dir / "labelmaps" baselines_dir = repo_root / "tests" / "baselines" # The workflow writes its dataset JSON, YAML config, and checkpoint tree @@ -92,10 +98,12 @@ if test_mode: data_dir = repo_root / "data" / "test" / "DirLab-4DCT" number_of_iterations_greedy: Optional[list[int]] = [1, 0] + number_of_iterations_icon = 1 epochs = 1 else: data_dir = repo_root / "data" / "DirLab-4DCT" - number_of_iterations_greedy = [60, 30, 20] # Greedy defaults + number_of_iterations_greedy = [60, 30, 20] + number_of_iterations_icon = 10 # 90 training frames at batch_size 4 is 22 optimizer steps per epoch, so # 100 epochs is ~2200 steps at a 5e-5 learning rate. Far fewer than # that leaves the finetuned weights statistically indistinguishable @@ -105,7 +113,7 @@ log_level = logging.INFO reporter = PhysioTwin4DBase(class_name=class_name, log_level=log_level) - output_dir.mkdir(parents=True, exist_ok=True) + labelmaps_dir.mkdir(parents=True, exist_ok=True) # Held-out evaluation pair (Case 1 is excluded from finetuning). T00 and # T50 are the extreme inhale/exhale phases, the only pair DIR-Lab supplies @@ -185,9 +193,18 @@ weights_path = workflow.process() else: weights_path = ( - Path(__file__).resolve().parent - / "network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/network_weights_final.trch" + experiment_dir + / f"{finetune_name}_model" + / "checkpoints" + / "network_weights_final.trch" ) + # Checked here rather than at the first set_weights_path() call, which + # only happens after the greedy and stock-ICON rows have already run. + if not weights_path.exists(): + raise FileNotFoundError( + f"run_finetuning is False but no checkpoint at {weights_path}. " + "Set run_finetuning = True to finetune from scratch." + ) # Registration comparison fixed_image = itk.imread(str(fixed_file), pixel_type=itk.F) @@ -235,8 +252,24 @@ def landmark_errors(transform: itk.Transform) -> np.ndarray: # transform, so Dice reflects the transform rather than what the segmenter # does differently on each interpolated volume. segmenter = SegmentNVSegmentCTMRI(log_level=log_level) - fixed_labelmap = segmenter.segment(fixed_image)["labelmap"] - moving_labelmap = segmenter.segment(moving_image)["labelmap"] + + def segment_phase(image_file: Path, image: itk.Image) -> itk.Image: + """Segment one phase, caching the labelmap under ``labelmaps_dir``. + + An existing labelmap short-circuits the segmentation, which dominates + this tutorial's runtime outside of finetuning. + """ + labelmap_file = labelmaps_dir / f"{image_file.stem}_labelmap.mha" + if labelmap_file.exists(): + reporter.log_info("Reusing cached labelmap: %s", labelmap_file.name) + return itk.imread(str(labelmap_file)) + + labelmap = segmenter.segment(image)["labelmap"] + itk.imwrite(labelmap, str(labelmap_file), compression=True) + return labelmap + + fixed_labelmap = segment_phase(fixed_file, fixed_image) + moving_labelmap = segment_phase(moving_file, moving_image) fixed_labels = itk.array_from_image(fixed_labelmap) def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: @@ -299,6 +332,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: ("greedy", None), ("icon_stock", None), ("icon_finetuned", weights_path), + ("greedy_icon_finetuned", weights_path), ): registrar: RegisterImagesBase if method_name == "greedy": @@ -306,12 +340,24 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: registrar.set_transform_type("Deformable") if number_of_iterations_greedy is not None: registrar.set_number_of_iterations(number_of_iterations_greedy) + elif method_name == "greedy_icon_finetuned": + # Both stages are configured exactly as the standalone "greedy" and + # "icon_finetuned" rows above, so this row differs from + # "icon_finetuned" only by the Greedy transform ICON starts from. + chain = RegisterImagesGreedyICON(log_level=log_level) + chain.greedy.set_transform_type("Deformable") + if number_of_iterations_greedy is not None: + chain.greedy.set_number_of_iterations(number_of_iterations_greedy) + chain.icon.set_number_of_iterations(number_of_iterations_icon) + chain.icon.set_mass_preservation(True) # For non-contrast CT + chain.icon.set_weights_path(str(method_weights)) + registrar = chain else: registrar = RegisterImagesICON(log_level=log_level) # None, not 0: icon_registration rejects 0 and takes None to mean # "no test-time finetuning steps", so the comparison reflects what # each set of weights predicts rather than per-pair optimization. - registrar.set_number_of_iterations(None) + registrar.set_number_of_iterations(number_of_iterations_icon) registrar.set_mass_preservation(True) # For non-contrast CT if method_weights is not None: registrar.set_weights_path(str(method_weights)) @@ -371,12 +417,12 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: len(fixed_landmarks), ) reporter.log_info( - " %-13s %7s %7s %7s %7s %9s", "method", "mean", "std", "p95", "max", "time_s" + " %-21s %7s %7s %7s %7s %9s", "method", "mean", "std", "p95", "max", "time_s" ) for row in rows: elapsed = row["registration_time_s"] reporter.log_info( - " %-13s %7.2f %7.2f %7.2f %7.2f %9s", + " %-21s %7.2f %7.2f %7.2f %7.2f %9s", row["method"], row["tre_mean"], row["tre_std"], @@ -387,7 +433,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: reporter.log_info("Per-class Dice of the warped moving labelmap against the fixed") reporter.log_info( - " %-13s %7s %7s %7s %7s %7s %7s %7s %12s", + " %-21s %7s %7s %7s %7s %7s %7s %7s %12s", "method", "classes", "mean", @@ -400,7 +446,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: ) for row in rows: reporter.log_info( - " %-13s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", + " %-21s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", row["method"], row["n_classes"], row["dice_mean"], diff --git a/tutorials/tutorial_04_heart_ct_to_vtk.py b/tutorials/tutorial_04_heart_ct_to_vtk.py index 069f0d9d..00f5cec9 100644 --- a/tutorials/tutorial_04_heart_ct_to_vtk.py +++ b/tutorials/tutorial_04_heart_ct_to_vtk.py @@ -110,9 +110,18 @@ ) # Result saving + # + # Merging the per-structure surfaces, rather than the per-group ones, lets + # the combined file carry a per-cell SegmentationLabelIds array: structure + # identity survives the merge, so the file can still be split per structure + # downstream. Per-group surfaces are contoured from a merged binary mask + # and have no per-cell identity to record. + combined_input = ( + result["label_surfaces"] if save_label_surfaces else result["surfaces"] + ) surface_file = Path( ContourTools.save_combined_surfaces( - result["surfaces"], + combined_input, str(output_dir / "patient_surfaces.vtp"), ) ) diff --git a/tutorials/tutorial_04_lung_ct_to_vtk.py b/tutorials/tutorial_04_lung_ct_to_vtk.py index 5b3260c5..be1738c6 100644 --- a/tutorials/tutorial_04_lung_ct_to_vtk.py +++ b/tutorials/tutorial_04_lung_ct_to_vtk.py @@ -94,9 +94,18 @@ ) # Result saving + # + # Merging the per-structure surfaces, rather than the per-group ones, lets + # the combined file carry a per-cell SegmentationLabelIds array: structure + # identity survives the merge, so the file can still be split per structure + # downstream. Per-group surfaces are contoured from a merged binary mask + # and have no per-cell identity to record. + combined_input = ( + result["label_surfaces"] if save_label_surfaces else result["surfaces"] + ) surface_file = Path( ContourTools.save_combined_surfaces( - result["surfaces"], + combined_input, str(output_dir / "patient_surfaces.vtp"), ) ) diff --git a/tutorials/tutorial_05_heart_vtk_to_usd.py b/tutorials/tutorial_05_heart_vtk_to_usd.py index 4f5d0fc4..327d42ff 100644 --- a/tutorials/tutorial_05_heart_vtk_to_usd.py +++ b/tutorials/tutorial_05_heart_vtk_to_usd.py @@ -6,9 +6,16 @@ Convert the VTK surface output from Tutorial 4, or another VTK-compatible mesh, into a USD file with anatomy materials. +Tutorial 4 writes one VTP per anatomical structure, each annotated with its +structure name. Feeding those files in individually — rather than the single +combined surface — keeps that name attached through the conversion, so each +structure becomes its own named USD prim and gets its own material: bright +oxygenated red for the left chambers, darker deoxygenated red for the right, +red-brown myocardium, and so on, instead of one uniform heart material. + Data Required ------------- -Preferred input: ``tutorials/output/tutorial_04_heart/patient_surfaces.vtp`` +Preferred input: ``tutorials/output/tutorial_04_heart/patient_*.vtp`` """ # Imports @@ -43,8 +50,8 @@ project_name = "tutorial_04_heart" - # Preferred input: the combined surface saved by Tutorial 4. - vtk_file = tutorials_dir / "output" / "tutorial_04_heart" / "patient_surfaces.vtp" + # Preferred input: the per-structure surfaces saved by Tutorial 4. + input_dir = tutorials_dir / "output" / "tutorial_04_heart" log_level = logging.INFO @@ -52,16 +59,42 @@ output_dir.mkdir(parents=True, exist_ok=True) - mesh = pv.read(str(vtk_file)) + # Tutorial 4 writes both per-group and per-structure surfaces into one + # directory. A per-structure surface is the one carrying exactly one name + # in SegmentationLabelNames, and its filename ends with that name; the + # per-group surfaces list every structure in the group and are skipped, as + # their geometry would otherwise be exported twice. + meshes: list[pv.PolyData] = [] + structure_names: list[str] = [] + for vtk_file in sorted(input_dir.glob("patient_*.vtp")): + mesh = pv.read(str(vtk_file)) + label_names = mesh.field_data.get("SegmentationLabelNames") + if label_names is None or len(label_names) != 1: + continue + if not vtk_file.stem.endswith(str(label_names[0])): + continue + meshes.append(mesh) + structure_names.append(str(label_names[0])) + + if not meshes: + raise FileNotFoundError( + "No per-structure surfaces found. Checked:\n" + + f" - {input_dir}/patient_*.vtp\n" + + "Run tutorial_04_heart_ct_to_vtk.py with save_label_surfaces=True." + ) # Workflow initialization - + # + # static_merge=True treats the meshes as separate objects in one scene + # rather than as frames of a time series. Leaving anatomy_type unset lets + # each object's name select its material, and object names default to the + # structure names read from field_data above. workflow = WorkflowConvertVTKToUSD( - input_meshes=[mesh], + input_meshes=meshes, usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", - anatomy_type="heart", + static_merge=True, separate_by_connectivity=True, log_level=log_level, ) @@ -84,4 +117,8 @@ ) ] - tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} + tutorial_results = { + "usd_file": results["usd_file"], + "structures": structure_names, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_06_heart_create_statistical_model.py b/tutorials/tutorial_06_heart_create_statistical_model.py index cf470bbc..b62a8499 100644 --- a/tutorials/tutorial_06_heart_create_statistical_model.py +++ b/tutorials/tutorial_06_heart_create_statistical_model.py @@ -48,10 +48,10 @@ test_mode = TestTools.running_as_test() if test_mode: data_dir = repo_root / "data" / "test" / "KCL-Heart-Model" - pca_components = 5 + number_of_pca_components = 5 else: data_dir = repo_root / "data" / "KCL-Heart-Model" - pca_components = 10 + number_of_pca_components = 10 log_level = logging.INFO @@ -87,7 +87,7 @@ workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_meshes, reference_mesh=reference_mesh, - pca_number_of_components=pca_components, + number_of_pca_components=number_of_pca_components, log_level=log_level, ) @@ -127,7 +127,7 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = min(2, pca_components, len(components), len(eigenvalues)) + mode_count = min(2, number_of_pca_components, len(components), len(eigenvalues)) xvfb_started = False try: diff --git a/tutorials/tutorial_06_lung_create_statistical_model.py b/tutorials/tutorial_06_lung_create_statistical_model.py index 6890491a..21ccd1b2 100644 --- a/tutorials/tutorial_06_lung_create_statistical_model.py +++ b/tutorials/tutorial_06_lung_create_statistical_model.py @@ -68,7 +68,7 @@ data_dir = repo_root / "data" / "DirLab-4DCT" - pca_number_of_modes = 6 + number_of_pca_components = 5 # Atlas iterations used to build the reference surface; 1 is a single # template-biased pass. @@ -128,7 +128,7 @@ workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_surfaces, reference_mesh=reference_surface, - pca_number_of_components=pca_number_of_modes, + number_of_pca_components=number_of_pca_components, log_level=log_level, ) @@ -167,7 +167,7 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = pca_number_of_modes + mode_count = number_of_pca_components mode_surface_files: list[Path] = [] xvfb_started = False diff --git a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py index 6438afe0..20614dd5 100644 --- a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py @@ -14,6 +14,10 @@ Patient image: a routine clinical 3D chest CT, ``data/Chest-CT/Chest-CT.mha``, downloaded with ``physiotwin4d-download-data Chest-CT --directory data/Chest-CT`` +ICON weights: ``tutorial_02_lung_distancemap_finetune_icon.py`` output +(``network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch``), +optional -- the stock uniGradICON weights are used when it is absent. """ # Imports @@ -57,10 +61,23 @@ pca_json = tutorial_06_dir / "pca_model.json" pca_mean_file = tutorial_06_dir / "pca_mean_surface.vtp" - pca_number_of_modes = 5 + number_of_pca_components = 5 patient_image_file = repo_root / "data" / "Chest-CT" / "Chest-CT.mha" + # Distance-map weights finetuned on DIR-Lab by Tutorial 2; see + # WorkflowFinetuneICONRegistration.expected_weights_path(). The + # mask_dilation_mm set below must match the one that tutorial finetuned + # with, since it fixes the distance maps' saturation radius. + icon_weights_path = ( + tutorials_dir + / "network_weights" + / "icon_dirlab_4dct_distancemap" + / "icon_dirlab_4dct_distancemap_model" + / "checkpoints" + / "network_weights_final.trch" + ) + log_level = logging.INFO # The same segmenter and surface-extraction workflow used by Tutorial 6, so @@ -124,12 +141,26 @@ workflow.set_use_pca_registration( use_pca_registration=True, pca_model=pca_model, - pca_number_of_modes=pca_number_of_modes, + number_of_pca_components=number_of_pca_components, use_surface=False, ) workflow.set_mask_dilation_mm(mask_dilation_mm=40) + # The labelmap-to-labelmap stage registers distance maps, not intensities, + # so it uses the distance-map-finetuned weights when they exist; without + # them the tutorial still runs, on the stock uniGradICON weights. + if icon_weights_path.exists(): + workflow.set_labelmap_to_labelmap_icon_weights_path(str(icon_weights_path)) + else: + workflow.log_warning( + "Finetuned distance-map ICON weights not found at %s; fitting with " + "the stock uniGradICON weights. Run " + "tutorials/tutorial_02_lung_distancemap_finetune_icon.py to create " + "them.", + icon_weights_path, + ) + # Workflow execution workflow_results = workflow.process() diff --git a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py index 176f4288..f40ec176 100644 --- a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py +++ b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py @@ -32,8 +32,11 @@ ``pca_mean_surface.vtp``) ICON weights: Tutorial 2 output (``network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/ -network_weights_final.trch``), optional — the stock uniGradICON weights are used -when it is absent. +network_weights_final.trch``) for the phase-to-phase CT registration, and +``network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch`` for +the distance-map stage of the SSM fit. Both optional — the stock uniGradICON +weights are used when they are absent. Outputs (per case, under ``output/tutorial_08_lung//``) ------------------------------------------------------------ @@ -101,6 +104,21 @@ / "network_weights_final.trch" ) + # Distance-map weights finetuned on DIR-Lab by + # tutorial_02_lung_distancemap_finetune_icon.py, used by the + # labelmap-to-labelmap stage of the SSM fit. mask_dilation_mm below must + # match the value that tutorial finetuned with, since it fixes the distance + # maps' saturation radius. + icon_distancemap_weights_path = ( + tutorials_dir + / "network_weights" + / "icon_dirlab_4dct_distancemap" + / "icon_dirlab_4dct_distancemap_model" + / "checkpoints" + / "network_weights_final.trch" + ) + fit_mask_dilation_mm = 40.0 + # Phase the SSM is fitted to; Tutorial 6 builds the lung PCA model from the # T70 surfaces, so the same phase is used here as the fitting reference. reference_phase = "T70" @@ -137,6 +155,21 @@ icon_weights_path, ) + use_finetuned_distancemap_weights = icon_distancemap_weights_path.exists() + if use_finetuned_distancemap_weights: + logger.info( + "Fitting the SSM with finetuned distance-map ICON weights: %s", + icon_distancemap_weights_path, + ) + else: + logger.warning( + "Finetuned distance-map ICON weights not found at %s; fitting the SSM " + "with the stock uniGradICON weights. Run " + "tutorials/tutorial_02_lung_distancemap_finetune_icon.py to create " + "them.", + icon_distancemap_weights_path, + ) + reference_files = sorted(data_dir.glob(f"Case*_{reference_phase}.mha")) if not reference_files: raise FileNotFoundError( @@ -200,9 +233,14 @@ fit_workflow.set_use_pca_registration( use_pca_registration=True, pca_model=pca_model, - pca_number_of_modes=6, + number_of_pca_components=6, use_surface=False, ) + fit_workflow.set_mask_dilation_mm(mask_dilation_mm=fit_mask_dilation_mm) + if use_finetuned_distancemap_weights: + fit_workflow.set_labelmap_to_labelmap_icon_weights_path( + str(icon_distancemap_weights_path) + ) fit_result = fit_workflow.process() pca_coefficients_file = case_output_dir / f"{case_id}_ssm_pca_coefficients.json" From 5b29e24d7d217ef294c350e02819ab77fe28cac8 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 13:54:10 -0400 Subject: [PATCH 4/5] ENH: Coderabbit --- .agents/agents/testing.md | 3 +- CLAUDE.md | 3 +- docs/tutorials.rst | 8 ++- .../simpleware_heart_segmentation.py | 5 ++ .../Reconstruct4DCT/reconstruct_4d_ct.py | 8 +-- src/physiotwin4d/convert_vtk_to_usd.py | 25 +++++-- src/physiotwin4d/register_images_ants.py | 11 +--- src/physiotwin4d/register_images_base.py | 6 ++ src/physiotwin4d/register_images_greedy.py | 2 +- .../register_models_distance_maps.py | 26 ++------ src/physiotwin4d/register_models_pca.py | 40 ++++++++--- .../register_time_series_images.py | 66 +++++++------------ src/physiotwin4d/transform_tools.py | 4 ++ .../workflow_convert_vtk_to_usd.py | 11 +++- ...rkflow_fit_statistical_model_to_patient.py | 36 +++++++--- statistics.md | 3 +- tests/README.md | 5 +- tests/conftest.py | 3 +- tests/test_contour_tools.py | 4 +- tests/test_register_models_pca.py | 6 ++ tests/test_workflow_convert_vtk_to_usd.py | 8 +-- tutorials/README.md | 2 +- ...torial_06_lung_create_statistical_model.py | 4 +- 23 files changed, 171 insertions(+), 118 deletions(-) diff --git a/.agents/agents/testing.md b/.agents/agents/testing.md index 83353ea5..9ee28736 100644 --- a/.agents/agents/testing.md +++ b/.agents/agents/testing.md @@ -29,7 +29,8 @@ python -m pytest tests/ -v # fast, recomm python -m pytest tests/test_contour_tools.py -v # single file python -m pytest tests/test_contour_tools.py::TestContourTools -v # single class python -m pytest tests/ -v --run-slow # opt into slow tests -python -m pytest tests/ -v --run-gpu --run-slow # typical local GPU profile (CI runner adds --run-simpleware --run-tutorials) +# typical local GPU profile; CI adds --run-simpleware --run-tutorials +python -m pytest tests/ -v --run-gpu --run-slow python -m pytest tests/ --create-baselines # create missing baselines ``` diff --git a/CLAUDE.md b/CLAUDE.md index 8ffc5fd3..5c13a0f7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -131,7 +131,8 @@ came from. tests are auto-skipped unless their opt-in flag is passed) py -m pytest tests/ -v - Baselines in `tests/baselines/` via Git LFS — run `git lfs pull` after cloning -- `tests/conftest.py`: session-scoped fixtures chaining download → convert → segment → register +- `tests/conftest.py`: session-scoped fixtures chaining + download → convert → segment → register - `src/physiotwin4d/test_tools.py`: baseline comparison utilities (`TestTools`, etc.) - Markers (all opt-in via `--run-`): `slow`, `requires_gpu`, `requires_simpleware`, `tutorial`. Data-dependent tests no diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 13a0b4c4..909ca591 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -11,7 +11,7 @@ Tutorials

PhysioTwin4D tutorials

From a CT scan to an animated digital twin

- Ten numbered stages across 15 runnable Python scripts. + Ten numbered stages across 16 runnable Python scripts. Each one drives the real workflow classes end-to-end on downloadable data, shows what it produced, and ends with the handful of constants to change so it runs on your own scans. @@ -219,6 +219,11 @@ Tutorial 2: Finetune ICON Registration Script ``tutorials/tutorial_02_lung_finetune_icon.py`` + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py`` — the + distance-map variant, which finetunes on distance maps rather than image + intensities so the labelmap-to-labelmap stage of Tutorials 7 and 8 has + in-distribution weights. + Workflow :class:`~physiotwin4d.WorkflowFinetuneICONRegistration`, then :class:`~physiotwin4d.RegisterImagesGreedy` and @@ -269,6 +274,7 @@ Run .. code-block:: bash python tutorials/tutorial_02_lung_finetune_icon.py + python tutorials/tutorial_02_lung_distancemap_finetune_icon.py Outputs The finetuned checkpoint under diff --git a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py index f1c9ba37..9ff7f4a0 100644 --- a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py +++ b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py @@ -23,6 +23,7 @@ # %% import logging import os +import sys import tkinter as tk from tkinter import filedialog @@ -66,6 +67,10 @@ ) root.destroy() +if not input_image_path: + print("No image selected; nothing to segment.") + sys.exit(0) + # Load the image try: input_image = itk.imread(input_image_path) diff --git a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py index c6c36cb4..a5c4ba90 100644 --- a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py +++ b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py @@ -162,9 +162,7 @@ def register_slices( # Try identity as initial transform print(" Trying init with identity.") - results_init_identity = reg_tool.register( - img, initial_forward_transform=None - ) + results_init_identity = reg_tool.register(img) inverse_tranform_init_identity = results_init_identity["inverse_transform"] forward_transform_init_identity = results_init_identity["forward_transform"] loss_init_identity = results_init_identity["loss"] @@ -173,8 +171,8 @@ def register_slices( if portion_of_prior_to_use > 0.0: # Try with prior transform print(" Trying with init prior.") - results_init_prior = reg_tool.register( - img, initial_forward_transform=prior_forward_transform + results_init_prior = reg_tool.register_from( + prior_forward_transform, img ) inverse_transform_init_prior = results_init_prior["inverse_transform"] forward_transform_init_prior = results_init_prior["forward_transform"] diff --git a/src/physiotwin4d/convert_vtk_to_usd.py b/src/physiotwin4d/convert_vtk_to_usd.py index b6df3216..6ccc1831 100644 --- a/src/physiotwin4d/convert_vtk_to_usd.py +++ b/src/physiotwin4d/convert_vtk_to_usd.py @@ -14,6 +14,7 @@ from __future__ import annotations import logging +from collections import Counter from collections.abc import Sequence from pathlib import Path from typing import Any, Literal, Optional, Union, cast @@ -165,11 +166,23 @@ def __init__( "time_codes must be in non-decreasing order; " "got values that decrease between consecutive frames" ) - if object_names is not None and len(object_names) != len(self.input_polydata): - raise ValueError( - f"object_names length ({len(object_names)}) must match " - f"input_polydata length ({len(self.input_polydata)})" - ) + if object_names is not None: + if len(object_names) != len(self.input_polydata): + raise ValueError( + f"object_names length ({len(object_names)}) must match " + f"input_polydata length ({len(self.input_polydata)})" + ) + # Each name becomes a prim path component, so it has to be a legal + # USD identifier and unique or prims silently collide. + invalid = [n for n in object_names if not Sdf.Path.IsValidIdentifier(n)] + counts = Counter(object_names) + duplicated = sorted(n for n, count in counts.items() if count > 1) + if invalid or duplicated: + raise ValueError( + "object_names must be unique and valid USD prim names " + "(letter or underscore followed by letters, digits or " + f"underscores); invalid: {invalid}, duplicated: {duplicated}" + ) self._is_static_merge: bool = static_merge self._time_codes: Optional[list[float]] = time_codes self.object_names: Optional[list[str]] = ( @@ -954,7 +967,7 @@ def _split_by_labels( elif "boundary_labels" in vtk_mesh.cell_data: label_array = vtk_mesh.cell_data["boundary_labels"] else: - self.logger.warning( + self.log_warning( "No 'SegmentationLabelIds' or 'boundary_labels' array found " "- using unified mesh" ) diff --git a/src/physiotwin4d/register_images_ants.py b/src/physiotwin4d/register_images_ants.py index a19df20b..cd32134d 100644 --- a/src/physiotwin4d/register_images_ants.py +++ b/src/physiotwin4d/register_images_ants.py @@ -539,20 +539,15 @@ def registration_method( consistent. The forward and inverse transforms are stored separately by ANTs. - IMPORTANT: the initial transform is applied by pre-warping the - moving image onto the fixed grid (the same approach as - RegisterImagesICON) rather than via ants.registration's - initial_transform argument, which mishandles matrix (affine/ - translation) initials. This method composes the initial transform - with the registration result, so the returned transforms include - both the initial alignment and the registration refinement. + To seed the registration with a known alignment, use + RegisterImagesBase.register_from(), which handles the pre-warp and + the composition. Implementation details: - Uses ANTs registration with configurable transform types - Supports multi-resolution optimization - Handles masked and unmasked registration - Returns ITK-compatible displacement field transforms - - Initial transforms are converted from ITK to ANTs format automatically Example: >>> # Basic registration diff --git a/src/physiotwin4d/register_images_base.py b/src/physiotwin4d/register_images_base.py index fba61d2d..7f3c1971 100644 --- a/src/physiotwin4d/register_images_base.py +++ b/src/physiotwin4d/register_images_base.py @@ -405,6 +405,12 @@ def register_from( initial_forward_transform, result, moving_image ) + # register() left the pre-warped image on self; the composed transforms + # are defined against the original, so restore it and drop any + # registered-image cache built for the pre-warped one. + self.moving_image = moving_image + self.moving_image_registered = None + self.forward_transform = composed["forward_transform"] self.inverse_transform = composed["inverse_transform"] self.loss = composed["loss"] diff --git a/src/physiotwin4d/register_images_greedy.py b/src/physiotwin4d/register_images_greedy.py index cd710c44..bf745ce0 100644 --- a/src/physiotwin4d/register_images_greedy.py +++ b/src/physiotwin4d/register_images_greedy.py @@ -363,7 +363,7 @@ def _registration_method_deformable( if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: cmd_aff += " -w 0.60" cmd_aff += " -i fixed moving" - kwargs_aff = { + kwargs_aff: dict[str, Any] = { "fixed": fixed_sitk, "moving": moving_sitk, } diff --git a/src/physiotwin4d/register_models_distance_maps.py b/src/physiotwin4d/register_models_distance_maps.py index 41867d22..43998b54 100644 --- a/src/physiotwin4d/register_models_distance_maps.py +++ b/src/physiotwin4d/register_models_distance_maps.py @@ -134,6 +134,8 @@ def __init__( 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. + distance_squared_max: Maximum squared distance, in squared millimeters, + that the distance maps are normalized against. Default: 50.0 mask_dilation_mm: Dilation amount in millimeters for binary registration mask generation. Default: 20mm log_level: Logging level (default: logging.INFO) @@ -227,15 +229,6 @@ def _create_masks_from_models(self) -> None: else: self.fixed_mask_image = None - itk.imwrite( - self.fixed_mask_image, "debug_fixed_mask_image.nii.gz", compression=True - ) - itk.imwrite( - self.fixed_distance_map_image, - "debug_fixed_distance_map_image.nii.gz", - compression=True, - ) - # Create moving distance map self.moving_distance_map_image = self.contour_tools.create_distance_map( self.moving_model, @@ -268,15 +261,6 @@ def _create_masks_from_models(self) -> None: else: self.moving_mask_image = None - itk.imwrite( - self.moving_mask_image, "debug_moving_mask_image.nii.gz", compression=True - ) - itk.imwrite( - self.moving_distance_map_image, - "debug_moving_distance_map_image.nii.gz", - compression=True, - ) - self.log_info("Distance map and mask generation complete") def register( @@ -398,18 +382,20 @@ def register( # (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). + # combine_displacement_field_transforms(a, b) evaluates b then a, so + # the stage that runs first is the second argument. self.forward_transform = ( self.transform_tools.combine_displacement_field_transforms( - forward_transform_ICON, forward_transform_Greedy, + forward_transform_ICON, reference_image=self.reference_image, mode="compose", ) ) self.inverse_transform = ( self.transform_tools.combine_displacement_field_transforms( - inverse_transform_Greedy, inverse_transform_ICON, + inverse_transform_Greedy, reference_image=self.reference_image, mode="compose", ) diff --git a/src/physiotwin4d/register_models_pca.py b/src/physiotwin4d/register_models_pca.py index c2cf68d0..fd42293a 100644 --- a/src/physiotwin4d/register_models_pca.py +++ b/src/physiotwin4d/register_models_pca.py @@ -217,6 +217,10 @@ def __init__( self.pca_template_model_point_subsample = pca_template_model_point_subsample self.pca_prior_weight = pca_prior_weight + if not 0.0 <= symmetric_weight <= 1.0: + raise ValueError( + f"symmetric_weight must be in [0, 1]; got {symmetric_weight}" + ) self.symmetric_weight = symmetric_weight # outputs @@ -479,12 +483,24 @@ def apply(vector: np.ndarray) -> np.ndarray: result = transform.TransformPoint(point) return np.array([result[0], result[1], result[2]], dtype=np.float64) - offset = apply(np.zeros(3)) + # Probe inside the model's own extent: a displacement field evaluated + # outside its grid returns no displacement, so probing at the origin and + # the unit cube would report such a transform as the identity affine. + bounds = np.asarray(self.pca_template_model.bounds, dtype=np.float64) + low, high = bounds[0::2], bounds[1::2] + center = 0.5 * (low + high) + step = 0.25 * np.maximum(high - low, 1.0) + + base = apply(center) matrix = np.column_stack( - [apply(basis) - offset for basis in np.eye(3, dtype=np.float64)] + [ + (apply(center + step[i] * basis) - base) / step[i] + for i, basis in enumerate(np.eye(3, dtype=np.float64)) + ] ) - probe = np.array([0.37, -0.61, 0.83], dtype=np.float64) - scale = max(1.0, float(np.abs(matrix).max()), float(np.abs(offset).max())) + offset = base - matrix @ center + probe = center + step * np.array([0.37, -0.61, 0.83], dtype=np.float64) + scale = max(1.0, float(np.abs(base).max()), float(np.abs(matrix).max())) if not np.allclose(apply(probe), matrix @ probe + offset, atol=1e-9 * scale): return None return matrix, offset @@ -1011,20 +1027,26 @@ def _log_transform_fidelity(self, template_points: np.ndarray) -> None: "PCA deformation must be computed" ) + # Two TransformPoint calls per point is costly on dense templates, and a + # strided subset reports the same RMS to within sampling noise. + stride = max(1, len(template_points) // 5000) + sampled = template_points[::stride] + deformation = self.registered_model_pca_deformation[::stride] + point = itk.Point[itk.D, 3]() - forward = np.empty_like(template_points) - round_trip = np.empty_like(template_points) - for i, source in enumerate(template_points): + forward = np.empty_like(sampled) + round_trip = np.empty_like(sampled) + for i, source in enumerate(sampled): point[0], point[1], point[2] = (float(v) for v in source) mapped = self.forward_point_transform.TransformPoint(point) forward[i] = (mapped[0], mapped[1], mapped[2]) back = self.inverse_point_transform.TransformPoint(mapped) round_trip[i] = (back[0], back[1], back[2]) - expected = template_points + self.registered_model_pca_deformation + expected = sampled + deformation field_rms = float(np.sqrt(np.mean(np.sum((forward - expected) ** 2, axis=1)))) inverse_rms = float( - np.sqrt(np.mean(np.sum((round_trip - template_points) ** 2, axis=1))) + np.sqrt(np.mean(np.sum((round_trip - sampled) ** 2, axis=1))) ) self.log_info( "Deformation field RMS error: %.4f mm (approximation of the " diff --git a/src/physiotwin4d/register_time_series_images.py b/src/physiotwin4d/register_time_series_images.py index 3462588d..4da6f2e7 100644 --- a/src/physiotwin4d/register_time_series_images.py +++ b/src/physiotwin4d/register_time_series_images.py @@ -23,18 +23,9 @@ class RegisterTimeSeriesImages(RegisterImagesBase): """Register a time series of images to a fixed image. - This class extends RegisterImagesBase to provide sequential registration - of multiple images (time series) to a fixed image, using a - caller-supplied registration backend. It can propagate information from - prior registrations to initialize subsequent ones. - - The registration proceeds in two passes from a reference frame: - - 1. Forward pass: from reference_frame to the end of the series - 2. Backward pass: from reference_frame-1 to the beginning - - This bidirectional approach helps maintain temporal coherence in the - registration results. + This class extends RegisterImagesBase to provide registration of multiple + images (time series) to a fixed image, using a caller-supplied registration + backend. Every frame is registered to the fixed image independently. Key features: @@ -165,12 +156,8 @@ def register_time_series( """Register a time series of images to the fixed image. This method registers an ordered sequence of images to a common fixed - frame. Registration proceeds bidirectionally from a reference frame: - forward to the end and backward to the beginning. - - For each image after the reference image, the method can optionally use - the transform from the previous image to initialize the registration, - which can improve convergence and temporal coherence. + frame. The reference frame is registered first, then every other frame, + each independently of the others. Args: moving_images (list[itk.Image]): List of 3D images to register @@ -180,9 +167,8 @@ def register_time_series( moving_labelmaps (list[itk.Image], optional): Per-frame multi-label segmentations, one for each moving image. If None, no labelmaps are used. If provided, must have the same length as moving_images. Default: None - reference_frame (int, optional): Index of the reference image to register first. - Registration proceeds forward from this index to the end, then - backward from this index to the beginning. Default: 0 + reference_frame (int, optional): Index of the reference image, which + is registered first. Default: 0 register_reference (bool, optional): If True, register the reference image to the fixed image. If False, use identity transform for the reference image. Default: True @@ -303,29 +289,25 @@ def register_time_series( inverse_transforms[reference_frame] = inverse_transform losses[reference_frame] = loss - # Register forward and backward from reference frame - for step, start_idx, end_idx in [ - (1, reference_frame + 1, num_images), # Forward pass - (-1, reference_frame - 1, -1), # Backward pass - ]: - for img_idx in range(start_idx, end_idx, step): - moving_image = moving_images[img_idx] - moving_mask = ( - moving_masks[img_idx] if moving_masks is not None else None - ) - moving_labelmap = ( - moving_labelmaps[img_idx] if moving_labelmaps is not None else None - ) + # Register every remaining frame; each is independent of the others. + for img_idx in range(num_images): + if img_idx == reference_frame: + continue + moving_image = moving_images[img_idx] + moving_mask = moving_masks[img_idx] if moving_masks is not None else None + moving_labelmap = ( + moving_labelmaps[img_idx] if moving_labelmaps is not None else None + ) - result = self.registrar.register( - moving_image=moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, - ) + result = self.registrar.register( + moving_image=moving_image, + moving_mask=moving_mask, + moving_labelmap=moving_labelmap, + ) - forward_transforms[img_idx] = result["forward_transform"] - inverse_transforms[img_idx] = result["inverse_transform"] - losses[img_idx] = cast(float, result["loss"]) + forward_transforms[img_idx] = result["forward_transform"] + inverse_transforms[img_idx] = result["inverse_transform"] + losses[img_idx] = cast(float, result["loss"]) assert all(t is not None for t in forward_transforms) assert all(t is not None for t in inverse_transforms) diff --git a/src/physiotwin4d/transform_tools.py b/src/physiotwin4d/transform_tools.py index d6e119fb..b891156a 100644 --- a/src/physiotwin4d/transform_tools.py +++ b/src/physiotwin4d/transform_tools.py @@ -82,6 +82,10 @@ def combine_displacement_field_transforms( In ``add`` mode, returns a single displacement field transform with weighted summed vectors. In ``compose`` mode, returns a composite transform containing both weighted displacement field transforms. + + ``compose`` follows ITK's CompositeTransform convention, where the + last-added transform is applied first: the result evaluates + ``tfm1(tfm2(x))``, so ``tfm2`` is the stage that runs first. """ assert mode in ["add", "compose"], "Invalid mode" diff --git a/src/physiotwin4d/workflow_convert_vtk_to_usd.py b/src/physiotwin4d/workflow_convert_vtk_to_usd.py index 684b0358..a920a7f5 100644 --- a/src/physiotwin4d/workflow_convert_vtk_to_usd.py +++ b/src/physiotwin4d/workflow_convert_vtk_to_usd.py @@ -136,6 +136,8 @@ def _read_object_annotations(self) -> list[tuple[Optional[str], Optional[str]]]: """ annotations: list[tuple[Optional[str], Optional[str]]] = [] for mesh in self.input_meshes: + if not isinstance(mesh, pv.DataSet) and isinstance(mesh, vtk.vtkDataSet): + mesh = pv.wrap(mesh) if not isinstance(mesh, pv.DataSet): annotations.append((None, None)) continue @@ -226,9 +228,14 @@ def process(self) -> dict[str, Any]: # Anatomy group per object name, used as the fallback when the name # itself matches no material (e.g. "rib_left_3" -> the bone group). + # Keyed by the prim names ConvertVTKToUSD will actually emit, which fall + # back to "_" when no object_names were derived. object_groups: dict[str, str] = {} - if object_names is not None: - for object_name, (_, group) in zip(object_names, annotations): + if self.static_merge: + group_keys = object_names or [ + f"{self.usd_project_name}_{index}" for index in range(len(annotations)) + ] + for object_name, (_, group) in zip(group_keys, annotations): if group is not None: object_groups[object_name] = group diff --git a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py index e185a76a..eb34181e 100644 --- a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py @@ -23,6 +23,7 @@ """ import logging +from pathlib import Path from typing import Any, Optional, cast import itk @@ -311,7 +312,12 @@ def set_labelmap_to_labelmap_icon_weights_path(self, weights_path: str) -> None: Args: weights_path: Path to an existing uniGradICON checkpoint. + + Raises: + FileNotFoundError: If weights_path does not exist. """ + if not Path(weights_path).exists(): + raise FileNotFoundError(f"ICON weights not found: {weights_path}") self.l2l_icon_weights_path = weights_path def set_use_pca_registration( @@ -585,25 +591,32 @@ def register_model_to_model_pca(self) -> dict: algorithm="dataset_surface" ) + # The PCA field is splatted at the *un-aligned* template's points, which + # generally fall outside the patient image, so grid it in the template's + # own frame; create_deformation_field drops samples that land off-grid. + # The ICP alignment is applied separately, after this field. + pca_field_reference_image = self.contour_tools.create_reference_image( + pca_template_model, + spatial_resolution=float(min(self.patient_image.GetSpacing())), + ) pca_transforms = self.pca_registrar.compute_pca_transforms( - reference_image=self.patient_image, + reference_image=pca_field_reference_image, ) self.pca_forward_point_transform = pca_transforms["forward_point_transform"] self.pca_inverse_point_transform = pca_transforms["inverse_point_transform"] if self.log_level == logging.DEBUG: - tfm_arr = itk.GetArrayFromImage( - self.pca_forward_point_transform.GetDisplacementField() - ) + tfm_field = self.pca_forward_point_transform.GetDisplacementField() + tfm_arr = itk.GetArrayFromImage(tfm_field) tfm_x_arr = tfm_arr[:, :, :, 0] tfm_y_arr = tfm_arr[:, :, :, 1] tfm_z_arr = tfm_arr[:, :, :, 2] tfm_x_img = itk.GetImageFromArray(tfm_x_arr) tfm_y_img = itk.GetImageFromArray(tfm_y_arr) tfm_z_img = itk.GetImageFromArray(tfm_z_arr) - tfm_x_img.CopyInformation(self.patient_image) - tfm_y_img.CopyInformation(self.patient_image) - tfm_z_img.CopyInformation(self.patient_image) + tfm_x_img.CopyInformation(tfm_field) + tfm_y_img.CopyInformation(tfm_field) + tfm_z_img.CopyInformation(tfm_field) itk.imwrite(tfm_x_img, "pca_forward_point_transform_x.nii.gz") itk.imwrite(tfm_y_img, "pca_forward_point_transform_y.nii.gz") itk.imwrite(tfm_z_img, "pca_forward_point_transform_z.nii.gz") @@ -695,9 +708,14 @@ def register_labelmap_to_labelmap(self) -> Optional[dict]: # Create a padded patient image since often the surface of interest # is not fully contained within the original image, which causes trouble - # with distance map registration. + # with distance map registration. The margin is physical -- it has to + # hold the dilated masks -- so it is converted per axis rather than + # padding a fixed voxel count on grids of any spacing. + margin_mm = 2.5 * self.mask_dilation_mm + spacing = np.asarray(self.patient_image.GetSpacing(), dtype=np.float64) + pad_voxels = np.maximum(1, np.ceil(margin_mm / spacing)).astype(int).tolist() padded_patient_image = ImageTools().pad_image( - self.patient_image, pad_voxels=[50, 50, 50], background_value=-1000 + self.patient_image, pad_voxels=pad_voxels, background_value=-1000 ) labelmap_registrar = RegisterModelsDistanceMaps( moving_model=self.pca_template_model_surface, diff --git a/statistics.md b/statistics.md index 40f89c2a..a5a85941 100644 --- a/statistics.md +++ b/statistics.md @@ -152,7 +152,8 @@ PhysioTwin4D operates across several technically demanding domains: - `requires_gpu` - GPU/CUDA-dependent tests (opt-in via `--run-gpu`) - `requires_simpleware` - tests needing a local Synopsys Simpleware Medical install (opt-in via `--run-simpleware`) - `requires_physicsnemo` - tests needing the optional `[physicsnemo]` extra (opt-in via `--run-physicsnemo`) -- `tutorial` - runs tutorial scripts end-to-end (opt-in via `--run-tutorials`; multi-hour) +- `tutorial` - runs tutorial scripts end-to-end + (opt-in via `--run-tutorials`; multi-hour) --- diff --git a/tests/README.md b/tests/README.md index fe7c3a62..58eb2511 100644 --- a/tests/README.md +++ b/tests/README.md @@ -140,7 +140,10 @@ pytest tests/ --create-baselines ## Test Timing Reports -All test runs automatically generate a comprehensive timing report at the end showing individual test durations, session time, and pass/fail/skip counts. The report separates regular tests from tutorial tests and highlights the slowest tests. +All test runs automatically generate a comprehensive timing report at the end +showing individual test durations, session time, and pass/fail/skip counts. The +report separates regular tests from tutorial tests and highlights the slowest +tests. ## Test Configuration diff --git a/tests/conftest.py b/tests/conftest.py index 00af88a3..4befdafc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -616,8 +616,7 @@ class KnownShiftCase: ``shift_mm``, so ``moving(q) == fixed(q + shift_mm)``. Warping ``moving`` back onto the fixed grid therefore requires a ``forward_transform`` of ``-shift_mm``, which gives an absolute accuracy target instead of the - "did it return something" checks that let a Greedy RAS/LPS sign error go - unnoticed. + "did it return something" checks that let a sign error go unnoticed. """ def __init__(self, fixed_image: itk.Image, shift_mm: tuple[float, float, float]): diff --git a/tests/test_contour_tools.py b/tests/test_contour_tools.py index 7505aa6b..c5f81b03 100644 --- a/tests/test_contour_tools.py +++ b/tests/test_contour_tools.py @@ -341,7 +341,9 @@ class TestSaveCombinedSurfaces: """Structure identity must survive merging into one file.""" @staticmethod - def _annotated_sphere(center: tuple[float, float, float], label_id: int) -> Any: + def _annotated_sphere( + center: tuple[float, float, float], label_id: int + ) -> pv.PolyData: surface = pv.Sphere(radius=1.0, center=center) surface.field_data["SegmentationLabelIds"] = np.array( [label_id], dtype=np.int32 diff --git a/tests/test_register_models_pca.py b/tests/test_register_models_pca.py index 2ce2fbcd..884bc103 100644 --- a/tests/test_register_models_pca.py +++ b/tests/test_register_models_pca.py @@ -327,3 +327,9 @@ def test_pca_transforms_round_trip() -> None: assert field_rms < 1.5 assert round_trip_rms < 1.0 + + # An identity field would satisfy both bounds above, so require that the + # forward transform actually moves the points a comparable distance. + displacement_rms = np.sqrt(np.mean(np.sum((mapped - template_points) ** 2, axis=1))) + expected_rms = np.sqrt(np.mean(np.sum((expected - template_points) ** 2, axis=1))) + assert displacement_rms > 0.5 * expected_rms diff --git a/tests/test_workflow_convert_vtk_to_usd.py b/tests/test_workflow_convert_vtk_to_usd.py index 57da7e3b..1fc6ce2a 100644 --- a/tests/test_workflow_convert_vtk_to_usd.py +++ b/tests/test_workflow_convert_vtk_to_usd.py @@ -4,10 +4,10 @@ """ from pathlib import Path +from typing import Optional import numpy as np import pyvista as pv -import pytest from pxr import Usd, UsdShade from physiotwin4d import WorkflowConvertVTKToUSD @@ -16,7 +16,7 @@ def _labeled_sphere( center: tuple[float, float, float], label_name: str, - group: str | None = None, + group: Optional[str] = None, ) -> pv.PolyData: """Return a sphere annotated the way WorkflowConvertImageToVTK annotates one.""" surface = pv.Sphere(radius=1.0, center=center, theta_resolution=8, phi_resolution=8) @@ -167,7 +167,3 @@ def test_unlabeled_meshes_keep_positional_names(self, tmp_path: Path) -> None: stage = Usd.Stage.Open(result["usd_file"]) assert stage.GetPrimAtPath("/World/scan/scan_0_object1").IsValid() assert stage.GetPrimAtPath("/World/scan/scan_1_object1").IsValid() - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tutorials/README.md b/tutorials/README.md index 48f9050d..02c37c7f 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -28,7 +28,7 @@ current working directory. | 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | | 1 | [tutorial_01_lung_gated_ct_to_usd.py](tutorial_01_lung_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Lung gated 4D CT (prepare first) | | 2 | [tutorial_02_lung_finetune_icon.py](tutorial_02_lung_finetune_icon.py) | `WorkflowFinetuneICONRegistration` | DirLab-4DCT (manual) | -| 2 | [tutorial_02_lung_distancemap_finetune_icon.py](tutorial_02_lung_distancemap_finetune_icon.py) | `WorkflowFinetuneICONRegistration` on lung distance maps | DirLab-4DCT (manual) | +| 2 | [distancemap variant](tutorial_02_lung_distancemap_finetune_icon.py) | `WorkflowFinetuneICONRegistration` on distance maps | DirLab-4DCT (manual) | | 3 | [tutorial_03_heart_reconstruct_highres_4d_ct.py](tutorial_03_heart_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | Slicer-Heart-CT (prepare first) | | 3 | [tutorial_03_lung_reconstruct_highres_4d_ct.py](tutorial_03_lung_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | | 4 | [tutorial_04_heart_ct_to_vtk.py](tutorial_04_heart_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | diff --git a/tutorials/tutorial_06_lung_create_statistical_model.py b/tutorials/tutorial_06_lung_create_statistical_model.py index 21ccd1b2..8c475e6b 100644 --- a/tutorials/tutorial_06_lung_create_statistical_model.py +++ b/tutorials/tutorial_06_lung_create_statistical_model.py @@ -167,7 +167,9 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = number_of_pca_components + # PCA rank is capped by the sample count, so the model can hold fewer + # components than requested. + mode_count = min(number_of_pca_components, len(components), len(eigenvalues)) mode_surface_files: list[Path] = [] xvfb_started = False From 716250d67b182422073677a6e9594e5367004c4e Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Fri, 7 Aug 2026 14:26:44 -0400 Subject: [PATCH 5/5] ENH: tutorials on finetuning - improved output details --- docs/tutorials.rst | 3 ++- ...rkflow_fit_statistical_model_to_patient.py | 3 --- ...orial_02_lung_distancemap_finetune_icon.py | 12 ++++++++-- tutorials/tutorial_02_lung_finetune_icon.py | 22 +++++++++++-------- 4 files changed, 25 insertions(+), 15 deletions(-) diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 909ca591..8c25dc3f 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -279,7 +279,8 @@ Run Outputs The finetuned checkpoint under ``tutorials/network_weights/icon_dirlab_4dct/``, plus - ``registration_summary.csv``, the registered images, the fixed and warped + ``registration_summary.csv``, the fixed-minus-registered difference images + (residual structure is what separates the methods), the fixed and warped labelmaps, and before/after screenshots in ``tutorials/output/tutorial_02_lung/``. diff --git a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py index eb34181e..30ec9e33 100644 --- a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py @@ -617,9 +617,6 @@ def register_model_to_model_pca(self) -> dict: tfm_x_img.CopyInformation(tfm_field) tfm_y_img.CopyInformation(tfm_field) tfm_z_img.CopyInformation(tfm_field) - itk.imwrite(tfm_x_img, "pca_forward_point_transform_x.nii.gz") - itk.imwrite(tfm_y_img, "pca_forward_point_transform_y.nii.gz") - itk.imwrite(tfm_z_img, "pca_forward_point_transform_z.nii.gz") if self.use_surface: # forward_point_transform excludes the post-PCA step and is defined diff --git a/tutorials/tutorial_02_lung_distancemap_finetune_icon.py b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py index 6bc22b1e..1d754140 100644 --- a/tutorials/tutorial_02_lung_distancemap_finetune_icon.py +++ b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py @@ -493,10 +493,18 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: itk.imwrite( fixed_distance_map, str(output_dir / "fixed_distance_map.mha"), compression=True ) + # Difference against the fixed distance map, not the resampled result: + # residual structure is what distinguishes the methods, and it is invisible + # in the registered distance maps themselves. + fixed_arr = itk.GetArrayFromImage(fixed_distance_map).astype(np.float32) for method_name, distance_map in registered_distance_maps.items(): + difference = itk.GetImageFromArray( + fixed_arr - itk.GetArrayFromImage(distance_map).astype(np.float32) + ) + difference.CopyInformation(fixed_distance_map) itk.imwrite( - distance_map, - str(output_dir / f"registered_distance_map_{method_name}.mha"), + difference, + str(output_dir / f"difference_distance_map_{method_name}.mha"), compression=True, ) for method_name, labelmap in labelmaps.items(): diff --git a/tutorials/tutorial_02_lung_finetune_icon.py b/tutorials/tutorial_02_lung_finetune_icon.py index 2574789f..89b13e8d 100644 --- a/tutorials/tutorial_02_lung_finetune_icon.py +++ b/tutorials/tutorial_02_lung_finetune_icon.py @@ -98,16 +98,12 @@ if test_mode: data_dir = repo_root / "data" / "test" / "DirLab-4DCT" number_of_iterations_greedy: Optional[list[int]] = [1, 0] - number_of_iterations_icon = 1 + number_of_iterations_icon = None # [1] epochs = 1 else: data_dir = repo_root / "data" / "DirLab-4DCT" - number_of_iterations_greedy = [60, 30, 20] - number_of_iterations_icon = 10 - # 90 training frames at batch_size 4 is 22 optimizer steps per epoch, so - # 100 epochs is ~2200 steps at a 5e-5 learning rate. Far fewer than - # that leaves the finetuned weights statistically indistinguishable - # from the stock weights they started from. + number_of_iterations_greedy = None # [60, 30, 20] + number_of_iterations_icon = None # [10] epochs = 100 log_level = logging.INFO @@ -392,10 +388,18 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: itk.imwrite( fixed_labelmap, str(output_dir / "fixed_labelmap.mha"), compression=True ) + # Difference against the fixed image, not the resampled result: residual + # structure is what distinguishes the methods, and it is invisible in the + # registered images themselves. + fixed_arr = itk.GetArrayFromImage(fixed_image).astype(np.float32) for method_name, image in registered_images.items(): + difference = itk.GetImageFromArray( + fixed_arr - itk.GetArrayFromImage(image).astype(np.float32) + ) + difference.CopyInformation(fixed_image) itk.imwrite( - image, - str(output_dir / f"registered_{method_name}.mha"), + difference, + str(output_dir / f"difference_{method_name}.mha"), compression=True, ) for method_name, labelmap in labelmaps.items():