diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 5d0c2cd..12d6fac 100644 --- a/src/physiotwin4d/contour_tools.py +++ b/src/physiotwin4d/contour_tools.py @@ -71,6 +71,42 @@ def extract_contours( return contours + @staticmethod + def smooth_and_decimate_surface( + surface: pv.PolyData, + decimation_reduction: float, + smoothing_iterations: int, + ) -> pv.PolyData: + """Optionally decimate then smooth a surface (no-op when disabled). + + Decimation uses ``decimate_pro`` on a triangulated copy; because + ``decimate_pro`` discards cell data, per-cell ``boundary_labels`` (needed + for anatomy splitting downstream) are transferred back onto the decimated + cells from their nearest original cell so anatomy materials still apply. + Smoothing uses non-shrinking Taubin smoothing, which only moves points and + therefore preserves cells and their labels. + + Args: + surface: Input surface. + decimation_reduction: Fraction of triangles to remove (0.0 disables). + smoothing_iterations: Taubin smoothing iterations (0 disables). + + Returns: + The decimated and/or smoothed surface. + """ + conditioned = surface + if decimation_reduction > 0.0: + original = conditioned + conditioned = conditioned.triangulate().decimate_pro(decimation_reduction) + if "boundary_labels" in original.cell_data: + nearest = original.find_closest_cell(conditioned.cell_centers().points) + conditioned.cell_data["boundary_labels"] = np.asarray( + original.cell_data["boundary_labels"] + )[nearest] + if smoothing_iterations > 0: + conditioned = conditioned.smooth_taubin(n_iter=smoothing_iterations) + return conditioned + @staticmethod def transform_contours( contours: pv.PolyData, diff --git a/src/physiotwin4d/transform_tools.py b/src/physiotwin4d/transform_tools.py index 401b427..75f86ef 100644 --- a/src/physiotwin4d/transform_tools.py +++ b/src/physiotwin4d/transform_tools.py @@ -12,23 +12,16 @@ """ import logging -from collections.abc import Sequence -from pathlib import Path -from typing import Type, TypeAlias, cast +from typing import Type, cast import itk import numpy as np import pyvista as pv import SimpleITK as sitk import vtk -from numpy.typing import NDArray -from pxr import Gf, Sdf, Usd, UsdGeom from .image_tools import ImageTools from .physiotwin4d_base import PhysioTwin4DBase -from .vtk_to_usd import add_framing_camera - -FloatArray: TypeAlias = NDArray[np.float32] | NDArray[np.float64] class TransformTools(PhysioTwin4DBase): @@ -559,6 +552,34 @@ def smooth_transform( return tfm_smooth + def smooth_deformation_field_transform( + self, field: itk.Image, sigma: float + ) -> itk.DisplacementFieldTransform: + """Wrap a deformation field as a Gaussian-smoothed field transform. + + The float vector ``field`` is converted to a double-precision vector + field, wrapped as a :class:`itk.DisplacementFieldTransform` and + Gaussian-smoothed by ``sigma`` (physical millimeters). Smoothing spreads + a thin surface-shell field into a continuous deformation (and attenuates + its peak magnitude). + + Args: + field (itk.Image): Input vector deformation field. + sigma (float): Standard deviation of the Gaussian smoothing kernel + in physical units (millimeters). + + Returns: + itk.DisplacementFieldTransform: Smoothed field transform. + """ + field_double = ImageTools().convert_array_to_image_of_vectors( + itk.array_from_image(field), reference_image=field, ptype=itk.D + ) + field_transform = itk.DisplacementFieldTransform[itk.D, 3].New() + field_transform.SetDisplacementField(field_double) + return self.smooth_transform( + field_transform, sigma=sigma, reference_image=field + ) + def combine_transforms_with_masks( self, transform1: itk.Transform, @@ -822,416 +843,3 @@ def convert_field_to_grid_visualization( grid_image_tfm = self.transform_image(grid_image, tfm, reference_image) return grid_image_tfm - - def convert_itk_transform_to_usd_visualization( - self, - tfm: itk.Transform, - reference_image: itk.image, - output_filename: str, - visualization_type: str = "arrows", - subsample_factor: int = 4, - arrow_scale: float = 1.0, - magnitude_threshold: float = 0.0, - ) -> str: - """ - Convert an ITK transform to a USD visualization for NVIDIA Omniverse. - - Creates a USD file containing either arrows or flow lines that visualize - the displacement field from the transform. Arrows show direction and - magnitude at sampled points, while flow lines show particle trajectories - through the deformation field. - - Args: - tfm (itk.Transform): Input ITK transform to visualize. Can be any ITK - transform type (Affine, BSpline, DisplacementField, Composite, etc.) - reference_image (itk.image): Defines the spatial grid for sampling - the displacement field (spacing, size, origin, direction) - output_filename (str): Path to output USD file (e.g., "deformation.usda") - visualization_type (str): Type of visualization - either "arrows" or - "flowlines". Default is "arrows". - subsample_factor (int): Subsample the displacement field by this factor - in each dimension to reduce primitive count. Default is 4 (64x fewer - points). Higher values = fewer primitives = better performance. - arrow_scale (float): Scale factor for arrow length. Default is 1.0. - Increase to make arrows longer, decrease to make them shorter. - magnitude_threshold (float): Only visualize displacements with magnitude - greater than this threshold (in mm). Default is 0.0 (show all). - Use to filter out small/negligible displacements. - - Returns: - str: Path to the created USD file - - Example: - >>> # Create arrow visualization of registration transform - >>> usd_file = transform_tools.convert_transform_to_usd_visualization( - ... registration_transform, - ... reference_ct, - ... 'deformation_arrows.usda', - ... visualization_type='arrows', - ... subsample_factor=8, # 512x fewer arrows - ... arrow_scale=2.0, # 2x longer arrows - ... magnitude_threshold=1.0, # Only show displacements > 1mm - ... ) - >>> - >>> # Create flow line visualization - >>> usd_file = transform_tools.convert_transform_to_usd_visualization( - ... registration_transform, - ... reference_ct, - ... 'deformation_flowlines.usda', - ... visualization_type='flowlines', - ... subsample_factor=4, - ... ) - - Note: - - The USD file can be opened in NVIDIA Omniverse for 3D visualization - - Arrows are colored by displacement magnitude (blue=low, red=high) - - Flowlines show particle paths through the deformation field - - Subsampling is critical for performance - dense fields can have millions - of points, creating too many primitives for interactive visualization - """ - # Generate ITK displacement field from transform - displacement_field = self.convert_transform_to_displacement_field( - tfm, reference_image - ) - - displacement_array = itk.GetArrayFromImage(displacement_field) - - # Get image size - size = displacement_field.GetLargestPossibleRegion().GetSize() - - # Subsample the field to reduce number of primitives - displacement_array = displacement_array[ - ::subsample_factor, ::subsample_factor, ::subsample_factor, : - ] - subsampled_size = [ - size[0] // subsample_factor, - size[1] // subsample_factor, - size[2] // subsample_factor, - ] - - # Remove any existing file and evict any stale in-memory USD layer. - # USD caches layers globally by identifier, so a prior call in the - # same Python session can block CreateNew even after the file is gone. - output_path = Path(output_filename) - if output_path.exists(): - output_path.unlink() - stale_layer = Sdf.Layer.Find(str(output_path)) - if stale_layer is not None: - stale_layer.Clear() - del stale_layer - - # Create USD stage - stage = Usd.Stage.CreateNew(output_filename) - - # Set up stage metadata - stage.SetMetadata("upAxis", "Y") - UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y) - - # Create root xform - # root_xform = UsdGeom.Xform.Define(stage, "/DeformationVisualization") - - if visualization_type == "arrows": - self._create_arrow_visualization( - stage, - displacement_array, - displacement_field, - subsampled_size, - subsample_factor, - arrow_scale, - magnitude_threshold, - ) - elif visualization_type == "flowlines": - self._create_flowline_visualization( - stage, - displacement_array, - displacement_field, - subsampled_size, - subsample_factor, - magnitude_threshold, - ) - else: - raise ValueError( - f"Invalid visualization_type: {visualization_type}. " - "Must be 'arrows' or 'flowlines'." - ) - - # Framing camera with tight near-clip for Omniverse Kit viewer ergonomics. - add_framing_camera(stage) - - # Save the stage - stage.Save() - self.log_info("Created USD visualization: %s", output_filename) - self.log_info(" Type: %s", visualization_type) - self.log_info(" Points: %d", np.prod(subsampled_size)) - self.log_info(" Subsample factor: %d", subsample_factor) - - return output_filename - - def _create_arrow_visualization( - self, - stage: Usd.Stage, - displacement_array: FloatArray, - displacement_field: itk.Image, - size: Sequence[int], - subsample_factor: int, - arrow_scale: float, - magnitude_threshold: float, - ) -> None: - """Create arrow-based visualization of displacement field.""" - # Iterate through all points in the subsampled field - arrow_count = 0 - for k in range(size[2]): - for j in range(size[1]): - for i in range(size[0]): - displacement = displacement_array[k, j, i, :] - - # Calculate magnitude - magnitude = np.linalg.norm(displacement) - - # Skip if below threshold - if magnitude < magnitude_threshold: - continue - - # Calculate index in original image space - index = [ - i * subsample_factor, - j * subsample_factor, - k * subsample_factor, - ] - - # Convert index to physical/world position using ITK - world_pos = displacement_field.TransformIndexToPhysicalPoint(index) - - # Create arrow at this position - arrow_path = f"/DeformationVisualization/Arrow_{arrow_count}" - self._create_arrow_prim( - stage, - arrow_path, - world_pos, - displacement, - float(magnitude), - arrow_scale, - ) - arrow_count += 1 - - self.log_info(" Created %d arrows", arrow_count) - - def _create_arrow_prim( - self, - stage: Usd.Stage, - prim_path: str, - position: Sequence[float], - displacement: FloatArray, - magnitude: float, - arrow_scale: float, - ) -> None: - """Create a single arrow primitive representing a displacement vector.""" - # Create xform for the arrow - arrow_xform = UsdGeom.Xform.Define(stage, prim_path) - - # Create a cone for the arrow (pointing in +Y direction by default) - cone = UsdGeom.Cone.Define(stage, f"{prim_path}/cone") - - # Set cone size based on displacement magnitude - arrow_length = float(magnitude * arrow_scale) - cone.GetHeightAttr().Set(arrow_length) - cone.GetRadiusAttr().Set(arrow_length * 0.1) # 10% of length - - # Calculate rotation to align arrow with displacement direction - if magnitude > 1e-6: # Avoid division by zero - # Normalize displacement to get direction - direction = displacement / magnitude - - # Default cone points in +Y, we want it to point along displacement - # Create rotation matrix to align +Y with displacement direction - up = np.array([0, 1, 0]) - rotation_axis = np.cross(up, direction) - rotation_axis_norm = np.linalg.norm(rotation_axis) - - if rotation_axis_norm > 1e-6: # Not parallel - rotation_axis = rotation_axis / rotation_axis_norm - cos_angle = np.dot(up, direction) - angle = np.arccos(np.clip(cos_angle, -1.0, 1.0)) - - # Convert axis-angle to quaternion - half_angle = angle / 2.0 - sin_half = np.sin(half_angle) - quat = Gf.Quatd( - np.cos(half_angle), # w - rotation_axis[0] * sin_half, # x - rotation_axis[1] * sin_half, # y - rotation_axis[2] * sin_half, # z - ) - - # Apply rotation (using Orient for quaternion) - arrow_xform.AddXformOp(UsdGeom.XformOp.TypeOrient).Set(quat) - - # Position the arrow at the base point + half displacement - # (so arrow starts at point and extends along displacement) - arrow_position = [ - position[0] + displacement[0] * 0.5, - position[1] + displacement[1] * 0.5, - position[2] + displacement[2] * 0.5, - ] - arrow_xform.AddXformOp(UsdGeom.XformOp.TypeTranslate).Set( - Gf.Vec3f( - float(arrow_position[0]), - float(arrow_position[1]), - float(arrow_position[2]), - ) - ) - - # Color based on magnitude (blue to red colormap) - # Normalize magnitude for coloring (assume max ~20mm displacement) - max_expected_displacement = 20.0 - normalized_mag = min(magnitude / max_expected_displacement, 1.0) - - # Blue (low) to red (high) colormap - color = Gf.Vec3f(float(normalized_mag), 0.0, float(1.0 - normalized_mag)) - cone.GetDisplayColorAttr().Set([color]) - - def _create_flowline_visualization( - self, - stage: Usd.Stage, - displacement_array: FloatArray, - displacement_field: itk.Image, - size: Sequence[int], - subsample_factor: int, - magnitude_threshold: float, - ) -> None: - """Create flow line visualization by tracing streamlines through displacement - field.""" - # Seed points - use a sparser grid for flowline seeds - seed_step = 2 # Every other point in the already-subsampled grid - flowline_count = 0 - - for k in range(0, size[2], seed_step): - for j in range(0, size[1], seed_step): - for i in range(0, size[0], seed_step): - # Get displacement at seed point - displacement = displacement_array[k, j, i, :] - magnitude = np.linalg.norm(displacement) - - # Skip if below threshold - if magnitude < magnitude_threshold: - continue - - # Calculate index in original image space - index = [ - i * subsample_factor, - j * subsample_factor, - k * subsample_factor, - ] - - # Convert index to physical/world position using ITK - seed_pos = displacement_field.TransformIndexToPhysicalPoint(index) - - # Trace streamline from this seed point - streamline_points = self._trace_streamline( - displacement_array, - displacement_field, - subsample_factor, - seed_pos, - max_steps=50, - step_size=0.5, - ) - - # Create USD curve for this streamline - if len(streamline_points) > 1: - curve_path = ( - f"/DeformationVisualization/Flowlines/Line_{flowline_count}" - ) - self._create_curve_prim(stage, curve_path, streamline_points) - flowline_count += 1 - - self.log_info(" Created %d flowlines", flowline_count) - - def _trace_streamline( - self, - displacement_array: FloatArray, - displacement_field: itk.Image, - subsample_factor: int, - seed_pos: Sequence[float], - max_steps: int, - step_size: float, - ) -> list[NDArray[np.float64]]: - """Trace a streamline through the displacement field using forward Euler - integration.""" - points = [np.array(seed_pos, dtype=float)] - current_pos = np.array(seed_pos, dtype=float) - - # Get original image size for bounds checking - original_size = displacement_field.GetLargestPossibleRegion().GetSize() - - for _ in range(max_steps): - # Convert world position to array indices using ITK - index = displacement_field.TransformPhysicalPointToIndex(tuple(current_pos)) - - # Convert to subsampled array indices - subsampled_indices = [ - index[0] // subsample_factor, - index[1] // subsample_factor, - index[2] // subsample_factor, - ] - - # Get subsampled array size - subsampled_size = [ - original_size[0] // subsample_factor, - original_size[1] // subsample_factor, - original_size[2] // subsample_factor, - ] - - # Check bounds in subsampled array (array is [k, j, i, :]) - if ( - subsampled_indices[0] < 0 - or subsampled_indices[0] >= subsampled_size[0] - or subsampled_indices[1] < 0 - or subsampled_indices[1] >= subsampled_size[1] - or subsampled_indices[2] < 0 - or subsampled_indices[2] >= subsampled_size[2] - ): - break - - # Get displacement at current position (array is [k, j, i, :]) - displacement = displacement_array[ - subsampled_indices[2], subsampled_indices[1], subsampled_indices[0], : - ] - magnitude = np.linalg.norm(displacement) - - # Stop if displacement is too small - if magnitude < 0.1: - break - - # Normalize displacement for integration - direction = displacement / (magnitude + 1e-10) - - # Take step along direction - current_pos = current_pos + direction * step_size - points.append(current_pos.copy()) - - return points - - def _create_curve_prim( - self, stage: Usd.Stage, prim_path: str, points: Sequence[NDArray[np.float64]] - ) -> UsdGeom.BasisCurves: - """Create a USD BasisCurves primitive for a streamline.""" - curve = UsdGeom.BasisCurves.Define(stage, prim_path) - - # Set curve properties - curve.GetTypeAttr().Set(UsdGeom.Tokens.linear) # Linear segments between points - curve.GetWrapAttr().Set(UsdGeom.Tokens.nonperiodic) # Open curve - - # Set points - points_array = [Gf.Vec3f(*p) for p in points] - curve.GetPointsAttr().Set(points_array) - - # Set curve vertex counts (one curve with N points) - curve.GetCurveVertexCountsAttr().Set([len(points)]) - - # Set width for visibility - curve.GetWidthsAttr().Set([0.5]) - - # Set color (cyan for flowlines) - curve.GetDisplayColorAttr().Set([Gf.Vec3f(0.0, 1.0, 1.0)]) - - return curve diff --git a/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py index 56c4b9e..15940df 100644 --- a/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py +++ b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py @@ -38,12 +38,16 @@ from pathlib import Path from typing import Any, Optional -from physiotwin4d import WorkflowInferPhysicsNeMoMGN +import pyvista as pv + +from physiotwin4d import WorkflowConvertVTKToUSD, WorkflowInferPhysicsNeMoMGN TUTORIALS_DIR = Path(__file__).resolve().parent FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") # Tutorial 9 run directory to evaluate (matches that trainer's OUTPUT_DIR). -MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" +MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn_3" + +EPOCH = 1500 DEFAULT_SUBJECT = "pm0027" DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10_byod_mgn" / DEFAULT_SUBJECT @@ -92,12 +96,27 @@ def predict( """Predict cardiac surfaces for *subject* using the trained MeshGraphNet.""" manifest_path = _write_subject_manifest(subject, out_dir) infer = WorkflowInferPhysicsNeMoMGN(model_directory=MODEL_DIR, epoch=epoch) - return infer.predict(manifest_path, stages=stages, output_directory=out_dir) + result = infer.predict(manifest_path, stages=stages, output_directory=out_dir) + + # Assemble the ordered predicted surfaces into a 4D USD, colored with the heart + # anatomy material via USDAnatomyTools (appearance="anatomy"). + surfaces = [pv.read(str(path)) for path in result["predicted_surfaces"]] + usd_workflow = WorkflowConvertVTKToUSD( + input_meshes=surfaces, + usd_project_name=f"{subject}_mgn_4d", + output_directory=out_dir, + appearance="anatomy", + anatomy_type="heart", + separate_by_connectivity=True, + frames_per_second=float(len(surfaces)), + ) + result["usd_file"] = usd_workflow.process()["usd_file"] + return result def run_tutorial() -> dict[str, Any]: """Tutorial / test entry point: evaluate DEFAULT_SUBJECT with the final weights.""" - return predict(DEFAULT_SUBJECT, DEFAULT_OUT_DIR) + return predict(DEFAULT_SUBJECT, DEFAULT_OUT_DIR, epoch=EPOCH) def main() -> None: diff --git a/tutorials/tutorial_11_heart_and_lung_motion.py b/tutorials/tutorial_11_heart_and_lung_motion.py index 385b427..1e3b780 100644 --- a/tutorials/tutorial_11_heart_and_lung_motion.py +++ b/tutorials/tutorial_11_heart_and_lung_motion.py @@ -61,6 +61,9 @@ heart alone. - ``combined_frame_.vtp`` (``000..099``) + ``heart_and_lung_motion.usd`` - the combined respiratory + cardiac 4D motion, painted with anatomy materials. +- ``combined_ct_.mha`` / ``combined_labelmap_.mha`` (``000..099``) - the + original CT and labelmap warped by the same per-frame combined deformation + (signed-short), so their anatomy tracks the displaced surfaces. Prerequisites ------------- @@ -71,9 +74,7 @@ from __future__ import annotations import logging -import shutil from pathlib import Path -from typing import cast import itk import numpy as np @@ -89,101 +90,6 @@ WorkflowConvertVTKToUSD, WorkflowInferPhysicsNeMoMGN, ) -from physiotwin4d import physicsnemo_tools as pnt - - -def _ensure_mgn_inference_assets( - model_dir: Path, epoch: int, pca_mean_volume: Path -) -> None: - """Complete an interrupted MGN run directory so it can be loaded for inference. - - ``WorkflowInferPhysicsNeMoMGN`` expects a finalized run directory - (``mgn_stage_model.pt``, ``pca_mean_surface.vtp`` and the shared graph - tensors). A run that only holds epoch checkpoints is completed here by - regenerating the missing assets deterministically: - - - ``mgn_stage_model.pt`` from the self-describing epoch checkpoint (it - carries the normalization stats and architecture the loader reads); - - ``pca_mean_surface.vtp`` and the shared MGN graph tensors from the PCA - template volume, using the same steps the trainer used. - - All writes are idempotent (skipped when the target already exists). - """ - import torch - - epoch_ckpt = model_dir / f"mgn_stage_model_epoch_{epoch:05d}.pt" - if not epoch_ckpt.exists(): - raise FileNotFoundError(f"Epoch checkpoint not found: {epoch_ckpt}") - - final_ckpt = model_dir / "mgn_stage_model.pt" - if not final_ckpt.exists(): - shutil.copy2(epoch_ckpt, final_ckpt) - - surface_file = model_dir / "pca_mean_surface.vtp" - if not surface_file.exists(): - volume = pv.read(str(pca_mean_volume)) - mean_surface = volume.extract_surface(algorithm="dataset_surface") - mean_surface.save(str(surface_file)) - mean_surface = cast(pv.PolyData, pv.read(str(surface_file))) - - edge_index_file = model_dir / "shared_edge_index.pt" - edge_feats_file = model_dir / "shared_edge_features.pt" - if not edge_index_file.exists() or not edge_feats_file.exists(): - edge_index = pnt.mesh_to_edge_index(mean_surface) - coords = np.asarray(mean_surface.points, dtype=np.float32) - edge_feats = pnt.compute_edge_features(coords, edge_index) - torch.save(edge_index, str(edge_index_file)) - torch.save(edge_feats, str(edge_feats_file)) - - -def _smoothed_cardiac_transform( - field: itk.Image, sigma_mm: float, transform_tools: TransformTools -) -> itk.DisplacementFieldTransform: - """Wrap a cardiac deformation field as a Gaussian-smoothed field transform. - - The float vector ``field`` is converted to a double-precision vector field, - wrapped as a ``DisplacementFieldTransform`` and Gaussian-smoothed by - ``sigma_mm`` (physical millimeters). Smoothing spreads the thin surface-shell - field into a continuous deformation (and attenuates its peak magnitude). - """ - field_double = ImageTools().convert_array_to_image_of_vectors( - itk.array_from_image(field), reference_image=field, ptype=itk.D - ) - field_transform = itk.DisplacementFieldTransform[itk.D, 3].New() - field_transform.SetDisplacementField(field_double) - return transform_tools.smooth_transform( - field_transform, sigma=sigma_mm, reference_image=field - ) - - -def _condition_surface( - surface: pv.PolyData, - decimation_reduction: float, - smoothing_iterations: int, -) -> pv.PolyData: - """Optionally decimate then smooth the model surface (no-op when disabled). - - Applied once to the labeled patient surface so every warped frame inherits - the same resolution and smoothing. Decimation uses ``decimate_pro`` on a - triangulated copy; because ``decimate_pro`` discards cell data, the per-cell - ``boundary_labels`` (needed for anatomy splitting downstream) are transferred - back onto the decimated cells from their nearest original cell so anatomy - materials still apply. Smoothing uses non-shrinking Taubin smoothing, which - only moves points and therefore preserves cells and their labels. - """ - conditioned = surface - if decimation_reduction > 0.0: - original = conditioned - conditioned = conditioned.triangulate().decimate_pro(decimation_reduction) - if "boundary_labels" in original.cell_data: - nearest = original.find_closest_cell(conditioned.cell_centers().points) - conditioned.cell_data["boundary_labels"] = np.asarray( - original.cell_data["boundary_labels"] - )[nearest] - if smoothing_iterations > 0: - conditioned = conditioned.smooth_taubin(n_iter=smoothing_iterations) - return conditioned - if __name__ == "__main__": logging.basicConfig(level=logging.INFO) @@ -193,7 +99,7 @@ def _condition_surface( tutorials_dir = Path(__file__).resolve().parent # ---- Hard-coded inputs (edit for your layout) -------------------------- - # Cardiac MGN model (Tutorial 9): epoch-300 checkpoint run directory and the + # Cardiac MGN model (Tutorial 9): epoch-1500 checkpoint run directory and the # PCA template volume the model was trained on. mgn_model_dir = tutorials_dir / "output" / "tutorial_09_byod_mgn_3" mgn_epoch = 1500 @@ -260,7 +166,6 @@ def _condition_surface( # ======================================================================== # Stage 1: cardiac motion - apply the MGN model to the Case1Pack heart. # ======================================================================== - _ensure_mgn_inference_assets(mgn_model_dir, mgn_epoch, pca_mean_volume) reference_image = itk.imread(str(reference_image_file)) infer = WorkflowInferPhysicsNeMoMGN(model_directory=mgn_model_dir, epoch=mgn_epoch) @@ -307,7 +212,7 @@ def _condition_surface( # Stage 2: combine the cardiac fields with respiratory motion. # ======================================================================== cardiac_transforms = [ - _smoothed_cardiac_transform(field, smoothing_sigma_mm, transform_tools) + transform_tools.smooth_deformation_field_transform(field, smoothing_sigma_mm) for field in cardiac_fields ] logger.info( @@ -322,7 +227,7 @@ def _condition_surface( contour_tools = ContourTools() patient_labelmap = itk.imread(str(patient_labelmap_file)) patient_surface = contour_tools.extract_contours(patient_labelmap) - patient_surface = _condition_surface( + patient_surface = contour_tools.smooth_and_decimate_surface( patient_surface, surface_decimation_reduction, surface_smoothing_iterations ) logger.info( @@ -342,11 +247,14 @@ def _condition_surface( n_phases = len(forward_transform_files) n_stages = len(cardiac_surfaces) + forward_transforms = [ + itk.transformread(str(forward_file)) for forward_file in forward_transform_files + ] + # Respiratory-warped vertex positions for every (phase, stage): # resp_points[phase][stage] = forward_phase(cardiac_surface[stage]).points. resp_points: list[list[np.ndarray]] = [] - for phase_idx, forward_file in enumerate(forward_transform_files): - forward_transform = itk.transformread(str(forward_file)) + for phase_idx, forward_transform in enumerate(forward_transforms): resp_points.append( [ np.asarray( @@ -360,8 +268,45 @@ def _condition_surface( ) logger.info("respiratory warp phase %d/%d done", phase_idx + 1, n_phases) - for stale_frame in output_dir.glob("combined_frame_*.vtp"): - stale_frame.unlink() + for pattern in ( + "combined_frame_*.vtp", + "combined_ct_*.mha", + "combined_labelmap_*.mha", + ): + for stale in output_dir.glob(pattern): + stale.unlink() + + # Combined displacement field d_ps(x) = forward_p(cardiac_s(x)) - x sampled on + # the reference grid, one per (phase, stage). Bilinearly blending these fields + # per frame reproduces the surface point-blend exactly (the blend is affine and + # the shared x cancels), so the warped CT/labelmap track the displaced surfaces. + # ITK CompositeTransform applies its last-added transform first, so cardiac is + # added last to act before respiratory. Fields are built lazily and evicted per + # phase so only the two phases bracketing the current frame are ever held. + image_tools = ImageTools() + field_cache: dict[tuple[int, int], np.ndarray] = {} + + def combined_field(phase: int, stage: int) -> np.ndarray: + cached = field_cache.get((phase, stage)) + if cached is not None: + return cached + forward = forward_transforms[phase] + composite = itk.CompositeTransform[itk.D, 3].New() + composite.AddTransform( + forward[0] if isinstance(forward, (list, tuple)) else forward + ) + composite.AddTransform(cardiac_transforms[stage]) + field = transform_tools.convert_transform_to_displacement_field( + composite, reference_image, np_component_type=np.float32 + ) + arr: np.ndarray = itk.array_from_image(field) + field_cache[(phase, stage)] = arr + return arr + + def to_signed_short(image: itk.Image) -> itk.Image: + caster = itk.CastImageFilter[type(image), itk.Image[itk.SS, 3]].New(Input=image) + caster.Update() + return caster.GetOutput() # Render frames by bilinearly interpolating the precomputed (phase, stage) # warped-point grid: respiratory advances with the breath phase, while the @@ -371,6 +316,8 @@ def _condition_surface( frames_per_phase = n_stages n_frames = n_phases * frames_per_phase combined_files: list[Path] = [] + ct_files: list[Path] = [] + labelmap_files: list[Path] = [] usd_frames: list[pv.PolyData] = [] for frame_idx in range(n_frames): # Respiratory position: current breath phase and fraction into it. @@ -403,9 +350,46 @@ def _condition_surface( combined_files.append(frame_file) usd_frames.append(combined_surface) + # Warp the original CT and labelmap by the same combined deformation, using + # the field bilinearly blended over the same four (phase, stage) corners. + field_a = (1.0 - card_blend) * combined_field(phase_idx, stage_idx) + ( + card_blend * combined_field(phase_idx, next_stage_idx) + ) + field_b = (1.0 - card_blend) * combined_field(next_phase_idx, stage_idx) + ( + card_blend * combined_field(next_phase_idx, next_stage_idx) + ) + field_arr = (1.0 - resp_blend) * field_a + resp_blend * field_b + field_img = image_tools.convert_array_to_image_of_vectors( + field_arr, ptype=itk.D, reference_image=reference_image + ) + field_transform = itk.DisplacementFieldTransform[itk.D, 3].New() + field_transform.SetDisplacementField(field_img) + + warped_ct = transform_tools.transform_image( + reference_image, field_transform, reference_image, "linear" + ) + warped_labelmap = transform_tools.transform_image( + patient_labelmap, field_transform, reference_image, "nearest" + ) + + ct_file = output_dir / f"combined_ct_{frame_idx:03d}.mha" + labelmap_file = output_dir / f"combined_labelmap_{frame_idx:03d}.mha" + itk.imwrite(warped_ct, str(ct_file), compression=True) + itk.imwrite( + to_signed_short(warped_labelmap), str(labelmap_file), compression=True + ) + ct_files.append(ct_file) + labelmap_files.append(labelmap_file) + + # Keep only the two phases bracketing the current frame in the field cache. + for key in [k for k in field_cache if k[0] not in (phase_idx, next_phase_idx)]: + del field_cache[key] + del resp_points logger.info( - "Wrote %d combined-motion surfaces to %s", len(combined_files), output_dir + "Wrote %d combined-motion surfaces, CT and labelmap volumes to %s", + len(combined_files), + output_dir, ) # Assemble the ordered frames into a single animated 4D USD, split by anatomy @@ -430,6 +414,8 @@ def _condition_surface( tutorial_results = { "cardiac_field_count": len(cardiac_fields), "combined_surfaces": combined_files, + "combined_ct_volumes": ct_files, + "combined_labelmap_volumes": labelmap_files, "beating_heart_usd": str(output_dir / "beating_heart.usd"), "usd_file": str(usd_file), }