From 59b563b8162974e6f1e517db91ef66d37e166916 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Thu, 16 Jul 2026 19:04:10 -0400 Subject: [PATCH 1/3] ENH: Demo for combining heart and lung motion --- docs/api/workflows.rst | 3 +- docs/tutorials.rst | 3 +- src/physiotwin4d/cli/convert_vtk_to_usd.py | 4 +- .../workflow_convert_vtk_to_usd.py | 11 +- .../workflow_infer_physicsnemo.py | 45 ++++- tutorials/tutorial_02_lung_ct_to_vtk.py | 156 ++++++++++++++++++ tutorials/tutorial_03_heart_vtk_to_usd.py | 42 ++--- ...o_lung_fit_statistical_model_to_patient.py | 48 +++--- .../tutorial_09_byod_train_physicsnemo_mgn.py | 2 +- 9 files changed, 251 insertions(+), 63 deletions(-) create mode 100644 tutorials/tutorial_02_lung_ct_to_vtk.py diff --git a/docs/api/workflows.rst b/docs/api/workflows.rst index 41502c5..1fc6227 100644 --- a/docs/api/workflows.rst +++ b/docs/api/workflows.rst @@ -119,7 +119,8 @@ VTK to USD anatomy_type="heart", ) - output_path = workflow.process() + result = workflow.process() + usd_file = result["usd_file"] Statistical Shape Modeling ========================== diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 9c5e594..6d752f6 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -368,7 +368,8 @@ Inner API usage anatomy_type="heart", separate_by_connectivity=True, ) - usd_file = workflow.process() + result = workflow.process() + usd_file = result["usd_file"] For callers who need more control than the workflow wrapper offers (e.g. applying a colormap or per-label anatomical splitting), use diff --git a/src/physiotwin4d/cli/convert_vtk_to_usd.py b/src/physiotwin4d/cli/convert_vtk_to_usd.py index 56feaea..a6d201b 100644 --- a/src/physiotwin4d/cli/convert_vtk_to_usd.py +++ b/src/physiotwin4d/cli/convert_vtk_to_usd.py @@ -230,9 +230,9 @@ def main() -> int: return 1 try: - out_path = workflow.process() + result = workflow.process() print("\nConversion completed successfully.") - print(f"Output: {out_path}") + print(f"Output: {result['usd_file']}") return 0 except Exception as e: print(f"\nError during conversion: {e}") diff --git a/src/physiotwin4d/workflow_convert_vtk_to_usd.py b/src/physiotwin4d/workflow_convert_vtk_to_usd.py index a5314fb..c3346ba 100644 --- a/src/physiotwin4d/workflow_convert_vtk_to_usd.py +++ b/src/physiotwin4d/workflow_convert_vtk_to_usd.py @@ -9,7 +9,7 @@ import logging from pathlib import Path -from typing import Literal, Optional, Sequence, Union +from typing import Any, Literal, Optional, Sequence, Union import pyvista as pv import vtk @@ -106,12 +106,13 @@ def __init__( "separate_by_connectivity and separate_by_cell_type cannot both be True" ) - def process(self) -> str: + def process(self) -> dict[str, Any]: """ Run the full workflow: convert meshes to USD, then apply the chosen appearance. Returns: - Path to the created USD file (str). + Dict with the results of the workflow: + - "usd_file" (str): Path to the created USD file. """ self.log_section("VTK to USD conversion workflow") @@ -168,7 +169,7 @@ def process(self) -> str: self.log_warning( "No mesh prims found under /World/%s", self.usd_project_name ) - return str(output_usd) + return {"usd_file": str(output_usd)} # Static merge has no time samples; pass None so only default time is used appearance_time_codes = None if self.static_merge else time_codes @@ -223,4 +224,4 @@ def process(self) -> str: primvar = None # next mesh: auto-pick again self.log_info("Workflow complete: %s", output_usd) - return str(output_usd) + return {"usd_file": str(output_usd)} diff --git a/src/physiotwin4d/workflow_infer_physicsnemo.py b/src/physiotwin4d/workflow_infer_physicsnemo.py index 019311b..3e64374 100644 --- a/src/physiotwin4d/workflow_infer_physicsnemo.py +++ b/src/physiotwin4d/workflow_infer_physicsnemo.py @@ -382,6 +382,7 @@ def create_deformation_field( stage: float, reference_image: itk.Image, output_directory: Optional[Path] = None, + reference_surface: Optional[Path] = None, ) -> dict[str, Any]: """Rasterize the inferred deformation onto a reference image grid. @@ -392,6 +393,15 @@ def create_deformation_field( (renormalized) reference-surface normal of those vertices. Empty voxels are zero. + By default the reference (undeformed) surface is reconstructed from the + PCA coefficients in the model's own frame. When the subject's reference + surface lives in a different world frame than the model template (e.g. a + patient scan whose statistical-model fit applied a pose transform not + captured by the shape coefficients), pass ``reference_surface`` so the + displacements are binned at the patient-space positions that actually + align with ``reference_image``. The network displacements themselves + depend only on the coefficients and stage, not on the binning positions. + Args: shape_parameters: JSON file with the subject PCA coefficient vector. stage: Target RR-interval fraction for the deformation. @@ -399,14 +409,34 @@ def create_deformation_field( (size, spacing, origin, direction). output_directory: If given, the two images are written there as compressed ``.mha`` files. + reference_surface: Optional mesh (volume or surface) whose extracted + surface supplies the binning positions and normals, overriding + the PCA reconstruction. Must share the mean-shape surface + topology (same point count and ordering); its surface is + extracted with the same ``dataset_surface`` algorithm used for + the model template, so a mesh built from the same PCA template + keeps the correspondence. Returns: Dict with ``deformation_field`` and ``normal_image`` (ITK vector - images) and, when written, their paths. + images), ``deformed_surface`` (the stage surface as ``pv.PolyData``) + and, when written, their paths. """ coeffs = pnt.load_pca_coefficients(shape_parameters) - mean_mesh, pca_model = self._load_pca_assets() - ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) + if reference_surface is not None: + patient_surface = cast( + pv.DataSet, pv.read(str(reference_surface)) + ).extract_surface(algorithm="dataset_surface") + ref_points = np.asarray(patient_surface.points, dtype=np.float32) + n_expected = len(self._mean_shape_coords) + if ref_points.shape[0] != n_expected: + raise ValueError( + f"reference_surface has {ref_points.shape[0]} surface points, " + f"expected {n_expected} (mean-shape topology)." + ) + else: + mean_mesh, pca_model = self._load_pca_assets() + ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) disps = self._predict_displacements(coeffs, stage) # Reference (undeformed) surface normals. @@ -452,19 +482,28 @@ def create_deformation_field( ref_points.shape[0], ) + # Deformed (stage) surface: reference positions displaced by the network, + # keeping the mean-shape topology. + deformed_surface = self._mean_surface.copy(deep=True) + deformed_surface.points = (ref_points + disps).astype(np.float32) + result: dict[str, Any] = { "deformation_field": deformation_image, "normal_image": normal_image, + "deformed_surface": deformed_surface, } if output_directory is not None: out_dir = Path(output_directory) out_dir.mkdir(parents=True, exist_ok=True) field_path = out_dir / "deformation_field.mha" normal_path = out_dir / "surface_normal_field.mha" + surface_path = out_dir / "deformed_surface.vtp" itk.imwrite(deformation_image, str(field_path), compression=True) itk.imwrite(normal_image, str(normal_path), compression=True) + deformed_surface.save(str(surface_path)) result["deformation_field_file"] = field_path result["normal_image_file"] = normal_path + result["deformed_surface_file"] = surface_path return result @staticmethod diff --git a/tutorials/tutorial_02_lung_ct_to_vtk.py b/tutorials/tutorial_02_lung_ct_to_vtk.py new file mode 100644 index 0000000..6256420 --- /dev/null +++ b/tutorials/tutorial_02_lung_ct_to_vtk.py @@ -0,0 +1,156 @@ +""" +Tutorial 2: CT Segmentation to VTK Surfaces + +Purpose +------- +Segment one 3D CT frame into anatomical groups and save a combined VTK +surface file. The output can be inspected directly in PyVista or used as +input for Tutorial 3. + +Data Required +------------- +Full data: ``data/DirLab-4DCT/Case1Pack_T??.mha`` +Test data: ``data/test/DirLab-4DCT/Case1Pack_T??.mha`` +""" + +# Imports +from __future__ import annotations + +import logging +from pathlib import Path + +import itk +import pyvista as pv + +from physiotwin4d import ( + ContourTools, + SegmentChestTotalSegmentator, + TestTools, + WorkflowConvertImageToVTK, +) + +# Only run if this script is not imported as a module + +# nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a +# multiprocessing.Pool. On Windows the spawn start method re-imports this +# script in each child; without the __name__ == "__main__" guard around +# top-level work, that re-import fires the segmenter again and Python's +# spawn-cascade detector raises RuntimeError. +if __name__ == "__main__": + # Data directory specification + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + project_name = "tutorial_02_lung" + + output_dir = tutorials_dir / "output" / project_name + + # In addition to the combined surface file always saved below, also + # save one VTP per anatomy group (e.g. heart.vtp, lung.vtp) and/or one + # VTP per individual anatomical structure (e.g. left_ventricle.vtp). + save_group_surfaces = True + save_label_surfaces = True + + test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + else: + data_dir = repo_root / "data" / "DirLab-4DCT" + + frame_files = sorted(data_dir.glob("Case1Pack_T??.mha")) + + log_level = logging.INFO + + segmentation_method = SegmentChestTotalSegmentator(log_level=log_level) + segmentation_method.set_has_academic_license(True) + + # Directory setup and data reading + output_dir.mkdir(parents=True, exist_ok=True) + + if not frame_files: + raise FileNotFoundError( + "DirLab-4DCT frame data not found. Checked:\n" + + f" - {data_dir}\n" + + "See data/README.md for download instructions." + ) + + ct_file = frame_files[0] + ct_image = itk.imread(str(ct_file)) + + # Workflow initialization + + workflow = WorkflowConvertImageToVTK( + segmentation_method=segmentation_method, + log_level=log_level, + ) + + # Workflow execution + # + # surface_target_reduction decimates each exported VTP surface. + result = workflow.process( + input_image=ct_image, + surface_target_reduction=0.5, + extract_label_surfaces=save_label_surfaces, + ) + + # Result saving + surface_file = Path( + ContourTools.save_combined_surface( + result["surfaces"], + str(output_dir), + prefix="patient", + ) + ) + if save_group_surfaces: + ContourTools.save_surfaces( + result["surfaces"], str(output_dir), prefix="patient" + ) + if save_label_surfaces: + ContourTools.save_surfaces( + result["label_surfaces"], str(output_dir), prefix="patient" + ) + labelmap_file = output_dir / "patient_labelmap.mha" + itk.imwrite(result["labelmap"], str(labelmap_file), compression=True) + + # Testing + tt = TestTools( + class_name=project_name, + results_dir=output_dir, + log_level=log_level, + ) + + screenshots: list[Path] = [] + screenshots.append( + tt.save_screenshot_image_slice( + ct_image, + f"{project_name}_segmentation_overlay.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + vmin=-200, + vmax=600, + overlay_mask=result["labelmap"], + ) + ) + + surfaces = [ + surface for surface in result["surfaces"].values() if surface is not None + ] + if surfaces: + combined_surface = pv.merge(surfaces) if len(surfaces) > 1 else surfaces[0] + screenshots.append( + tt.save_screenshot_mesh( + combined_surface, + f"{project_name}_vtk_surfaces.png", + camera_position="iso", + color="lightblue", + opacity=0.85, + ) + ) + + tutorial_results = { + "result": result, + "surface_file": surface_file, + "labelmap_file": labelmap_file, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_03_heart_vtk_to_usd.py b/tutorials/tutorial_03_heart_vtk_to_usd.py index 09063a1..750b194 100644 --- a/tutorials/tutorial_03_heart_vtk_to_usd.py +++ b/tutorials/tutorial_03_heart_vtk_to_usd.py @@ -9,7 +9,6 @@ Data Required ------------- Preferred input: ``tutorials/output/tutorial_02_heart/patient_surfaces.vtp`` -Fallback input: any ``*.vtp`` under ``data`` or ``data/test`` """ # Imports @@ -17,7 +16,6 @@ import logging from pathlib import Path -from typing import Optional import pyvista as pv @@ -43,18 +41,12 @@ output_dir = tutorials_dir / "output" / "tutorial_03_heart" baselines_dir = repo_root / "tests" / "baselines" - # Preferred input: the combined surface saved by Tutorial 2. Leave vtk_file as - # None to auto-discover (Tutorial 2 output first, then any *.vtp under data_dir). - tutorial_02_surface = ( + project_name = "tutorial_02_heart" + + # Preferred input: the combined surface saved by Tutorial 2. + vtk_file = ( tutorials_dir / "output" / "tutorial_02_heart" / "patient_surfaces.vtp" ) - vtk_file: Optional[Path] = None - - test_mode = TestTools.running_as_test() - if test_mode: - data_dir = repo_root / "data" / "test" - else: - data_dir = repo_root / "data" log_level = logging.INFO @@ -62,24 +54,13 @@ output_dir.mkdir(parents=True, exist_ok=True) - if vtk_file is None and tutorial_02_surface.exists(): - vtk_file = tutorial_02_surface - if vtk_file is None: - vtk_candidates = sorted(data_dir.rglob("*.vtp")) - if not vtk_candidates: - raise FileNotFoundError( - "No VTK surface file found. Run Tutorial 2 first or place a " - f"*.vtp file under {data_dir}." - ) - vtk_file = vtk_candidates[0] - mesh = pv.read(str(vtk_file)) # Workflow initialization workflow = WorkflowConvertVTKToUSD( input_meshes=[mesh], - usd_project_name="surfaces", + usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", anatomy_type="heart", @@ -88,7 +69,7 @@ ) # Workflow execution - usd_file = workflow.process() + results = workflow.process() # Testing tt = TestTools( @@ -98,12 +79,11 @@ log_level=log_level, ) - screenshots: list[Path] = [] - screenshots.append( + screenshots = [ tt.save_screenshot_openusd( - usd_file, - "usd_mesh_rendering.png", + results["usd_file"], + f"{project_name}_usd_mesh_rendering.png", ) - ) + ] - tutorial_results = {"usd_file": usd_file, "screenshots": screenshots} + tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} \ No newline at end of file diff --git a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py index 7354759..83a5993 100644 --- a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py @@ -45,9 +45,9 @@ repo_root = Path(__file__).resolve().parent.parent tutorials_dir = Path(__file__).resolve().parent - class_name = "tutorial_05_heart_to_lung_fit_statistical_model_to_patient" + project_name = "tutorial_05_heart_to_lung" - output_dir = tutorials_dir / "output" / "tutorial_05_heart_to_lung" + output_dir = tutorials_dir / "output" / project_name baselines_dir = repo_root / "tests" / "baselines" # PCA model + mean surface produced by Tutorial 4. @@ -55,6 +55,9 @@ pca_mean_file = ( tutorials_dir / "output" / "tutorial_04_heart" / "pca_mean_surface.vtp" ) + # BYOD example: + # pca_mean_file = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + # pca_json = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_model.json") test_mode = TestTools.running_as_test() if test_mode: @@ -94,18 +97,25 @@ "See data/README.md for download instructions." ) patient_image = itk.imread(str(patient_image_file)) - itk.imwrite(patient_image, output_dir / "patient_image.nii.gz") - segmentation_result = segmentation_method.segment(patient_image) - patient_labelmap = segmentation_result["labelmap"] - itk.imwrite(patient_labelmap, output_dir / "patient_labelmap.nii.gz") + if not (output_dir / f"{project_name}_patient_image.nii.gz").exists(): + itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz") - heart_labelmap = segmentation_result["heart"] - itk.imwrite(heart_labelmap, output_dir / "heart_labelmap.nii.gz") + segmentation_result = segmentation_method.segment(patient_image) + patient_labelmap = segmentation_result["labelmap"] + itk.imwrite(patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz") - contour_tools = ContourTools() - heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) - heart_surface.save(output_dir / "heart_surface.vtp") + heart_labelmap = segmentation_result["heart"] + itk.imwrite(heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz") + + contour_tools = ContourTools() + heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) + heart_surface.save(output_dir / f"{project_name}_heart_surface.vtp") + + else: + patient_labelmap = itk.imread(output_dir / f"{project_name}_patient_labelmap.nii.gz") + heart_labelmap = itk.imread(output_dir / f"{project_name}_heart_labelmap.nii.gz") + heart_surface = pv.read(output_dir / f"{project_name}_heart_surface.vtp") # Workflow initialization @@ -131,25 +141,25 @@ # Result saving registered_coefficients = workflow.pca_coefficients if registered_coefficients is not None: - registered_coefficients_path = output_dir / "registered_coefficients.json" + registered_coefficients_path = output_dir / f"{project_name}_registered_coefficients.json" with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: json.dump(registered_coefficients.tolist(), f) template_mesh = workflow.pca_template_model - template_mesh.save(str(output_dir / "template_mesh.vtp")) + template_mesh.save(str(output_dir / f"{project_name}_template_mesh.vtu")) template_surface = workflow.pca_template_model_surface - template_surface.save(str(output_dir / "template_surface.vtp")) + template_surface.save(str(output_dir / f"{project_name}_template_surface.vtp")) registered_mesh = workflow_results["registered_template_model"] - registered_mesh.save(str(output_dir / "template_mesh_registered.vtp")) + registered_mesh.save(str(output_dir / f"{project_name}_template_mesh_registered.vtu")) registered_surface = workflow_results["registered_template_model_surface"] - registered_surface.save(str(output_dir / "template_surface_registered.vtp")) + registered_surface.save(str(output_dir / f"{project_name}_template_surface_registered.vtp")) # Testing TestTools( - class_name=class_name, + class_name=project_name, results_dir=output_dir, baselines_dir=baselines_dir, log_level=log_level, @@ -162,7 +172,7 @@ screenshots: list[Path] = [] - before_path = output_dir / "model_before_registration.png" + before_path = output_dir / f"{project_name}_model_before_registration.png" plotter = pv.Plotter(off_screen=True, window_size=[800, 600]) plotter.add_mesh(pca_mean, color="dodgerblue", opacity=0.6) plotter.add_mesh(heart_surface, color="tomato", opacity=0.6) @@ -171,7 +181,7 @@ plotter.close() screenshots.append(before_path) - after_path = output_dir / "model_after_registration.png" + after_path = output_dir / f"{project_name}_model_after_registration.png" plotter = pv.Plotter(off_screen=True, window_size=[800, 600]) plotter.add_mesh(registered_surface, color="limegreen", opacity=0.7) plotter.add_mesh(heart_surface, color="tomato", opacity=0.4) diff --git a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py index ec664fe..a6276e5 100644 --- a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py @@ -90,7 +90,7 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" MANIFESTS_DIR = OUTPUT_DIR / "manifests_mgn" - RESUME_FROM = str(OUTPUT_DIR / "mgn_stage_model_epoch_00100.pt") + RESUME_FROM = str(TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn_2" / "mgn_stage_model_epoch_00200.pt") EPOCHS = 1500 BATCH_SIZE_GRAPHS = 4 # mini-batch measured in (subject, phase) graphs From 8f0ee747df27097d14f35b38c53b88d2976e810b Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sun, 19 Jul 2026 13:30:23 -0400 Subject: [PATCH 2/3] ENH: Combined heart+lung motion tutorial and per-organ USD materials Add Tutorial 11 that fuses cardiac motion (trained PhysicsNeMo MeshGraphNet applied to the patient heart) with respiratory motion (Tutorial 1 lung transforms) into one animated 4D USD, split by anatomy label and painted with per-organ OmniSurface materials. Add the Heart_and_Lungs_Motion experiment scripts behind it. usd_anatomy_tools: expand DEFAULT_RENDER_PARAMS with organ-level materials (skin, airway, vein/artery, muscle, fat, cartilage, endocrine organs, oxygenation-coded heart chambers) and switch enhance_meshes override matching to longest-substring-wins so e.g. "kidney" covers kidney_left/right. segment_chest_total_segmentator: move sacrum (label 25) into the correct group; stop overwriting body labels so body_skin (133) survives. Tutorials 05/09/10: evaluate the resumed run's output directory instead of OUTPUT_DIR, add missing casts/asserts/type hints for mypy. CI: run Ruff through pre-commit so its version is pinned in one place; bump ruff-pre-commit to v0.15.20. Add .claude-plugin marketplace/plugin manifests and .cursor rules; restructure CLAUDE.md behavior guidelines into numbered sections. --- .claude-plugin/marketplace.json | 29 + .claude-plugin/plugin.json | 11 + .cursor/rules/karpathy-guidelines.mdc | 70 +++ .cursor/rules/project-standards.mdc2 | 99 ++++ .github/workflows/ci.yml | 12 +- .gitignore | 1 - .pre-commit-config.yaml | 2 +- CLAUDE.md | 152 +++-- .../0-heart_and_lungs_beating_heart.py | 203 +++++++ .../1-heart_and_lungs_combined.py | 263 +++++++++ .../segment_chest_total_segmentator.py | 6 +- src/physiotwin4d/usd_anatomy_tools.py | 524 +++++++++++++++++- tutorials/tutorial_03_heart_vtk_to_usd.py | 6 +- ...o_lung_fit_statistical_model_to_patient.py | 38 +- .../tutorial_09_byod_train_physicsnemo_mgn.py | 23 +- .../tutorial_09_byod_train_physicsnemo_mlp.py | 9 +- .../tutorial_10_byod_eval_physicsnemo_mgn.py | 7 +- .../tutorial_10_byod_eval_physicsnemo_mlp.py | 7 +- .../tutorial_11_heart_and_lung_motion.py | 435 +++++++++++++++ 19 files changed, 1774 insertions(+), 123 deletions(-) create mode 100644 .claude-plugin/marketplace.json create mode 100644 .claude-plugin/plugin.json create mode 100644 .cursor/rules/karpathy-guidelines.mdc create mode 100644 .cursor/rules/project-standards.mdc2 create mode 100644 experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py create mode 100644 experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py create mode 100644 tutorials/tutorial_11_heart_and_lung_motion.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json new file mode 100644 index 0000000..f657354 --- /dev/null +++ b/.claude-plugin/marketplace.json @@ -0,0 +1,29 @@ +{ + "name": "karpathy-skills", + "id": "karpathy-skills", + "owner": { + "name": "forrestchang" + }, + "metadata": { + "description": "Behavioral guidelines to reduce common LLM coding mistakes, derived from Andrej Karpathy's observations", + "version": "1.0.0" + }, + "plugins": [ + { + "name": "andrej-karpathy-skills", + "source": "./", + "description": "Behavioral guidelines to reduce common LLM coding mistakes: Think Before Coding, Simplicity First, Surgical Changes, Goal-Driven Execution", + "version": "1.0.0", + "author": { + "name": "forrestchang" + }, + "keywords": [ + "guidelines", + "best-practices", + "coding", + "karpathy" + ], + "category": "workflow" + } + ] +} diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json new file mode 100644 index 0000000..6e88551 --- /dev/null +++ b/.claude-plugin/plugin.json @@ -0,0 +1,11 @@ +{ + "name": "andrej-karpathy-skills", + "description": "Behavioral guidelines to reduce common LLM coding mistakes, derived from Andrej Karpathy's observations on LLM coding pitfalls", + "version": "1.0.0", + "author": { + "name": "forrestchang" + }, + "license": "MIT", + "keywords": ["guidelines", "best-practices", "coding", "karpathy"], + "skills": ["./skills/karpathy-guidelines"] +} diff --git a/.cursor/rules/karpathy-guidelines.mdc b/.cursor/rules/karpathy-guidelines.mdc new file mode 100644 index 0000000..edd317f --- /dev/null +++ b/.cursor/rules/karpathy-guidelines.mdc @@ -0,0 +1,70 @@ +--- +description: Behavioral guidelines to reduce common LLM coding mistakes. Use when writing, reviewing, or refactoring code to avoid overcomplication, make surgical changes, surface assumptions, and define verifiable success criteria. +alwaysApply: true +--- + +# Karpathy behavioral guidelines + +Behavioral guidelines to reduce common LLM coding mistakes. Merge with project-specific instructions as needed. + +**Tradeoff:** These guidelines bias toward caution over speed. For trivial tasks, use judgment. + +## 1. Think Before Coding + +**Don't assume. Don't hide confusion. Surface tradeoffs.** + +Before implementing: +- State your assumptions explicitly. If uncertain, ask. +- If multiple interpretations exist, present them - don't pick silently. +- If a simpler approach exists, say so. Push back when warranted. +- If something is unclear, stop. Name what's confusing. Ask. + +## 2. Simplicity First + +**Minimum code that solves the problem. Nothing speculative.** + +- No features beyond what was asked. +- No abstractions for single-use code. +- No "flexibility" or "configurability" that wasn't requested. +- No error handling for impossible scenarios. +- If you write 200 lines and it could be 50, rewrite it. + +Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify. + +## 3. Surgical Changes + +**Touch only what you must. Clean up only your own mess.** + +When editing existing code: +- Don't "improve" adjacent code, comments, or formatting. +- Don't refactor things that aren't broken. +- Match existing style, even if you'd do it differently. +- If you notice unrelated dead code, mention it - don't delete it. + +When your changes create orphans: +- Remove imports/variables/functions that YOUR changes made unused. +- Don't remove pre-existing dead code unless asked. + +The test: Every changed line should trace directly to the user's request. + +## 4. Goal-Driven Execution + +**Define success criteria. Loop until verified.** + +Transform tasks into verifiable goals: +- "Add validation" → "Write tests for invalid inputs, then make them pass" +- "Fix the bug" → "Write a test that reproduces it, then make it pass" +- "Refactor X" → "Ensure tests pass before and after" + +For multi-step tasks, state a brief plan: +``` +1. [Step] → verify: [check] +2. [Step] → verify: [check] +3. [Step] → verify: [check] +``` + +Strong success criteria let you loop independently. Weak criteria ("make it work") require constant clarification. + +--- + +**These guidelines are working if:** fewer unnecessary changes in diffs, fewer rewrites due to overcomplication, and clarifying questions come before implementation rather than after mistakes. diff --git a/.cursor/rules/project-standards.mdc2 b/.cursor/rules/project-standards.mdc2 new file mode 100644 index 0000000..4413fa3 --- /dev/null +++ b/.cursor/rules/project-standards.mdc2 @@ -0,0 +1,99 @@ +--- +description: PhysioTwin4D project standards and workflow preferences +alwaysApply: true +--- + +# PhysioTwin4D Project Standards + +## File Operations + +**ALWAYS use git commands for file operations in this repository:** + +```bash +# Moving files +git mv old_path new_path + +# Deleting files +git rm file_path + +# Renaming files +git mv old_name.py new_name.py +``` + +❌ **Don't use**: `mv`, `rm`, `cp` directly +✅ **Do use**: `git mv`, `git rm` to maintain git history + +## Documentation + +**NEVER create .md files unless explicitly requested by the user or unless creating a new module in a directory that does not already have a README.md, then a README.md may be created if appropriate.** + +❌ **Don't create**: +- `MIGRATION.md` +- `CHANGES.md` +- `UPDATE_SUMMARY.md` +- `MODERNIZATION_*.md` +- `*_GUIDE.md` +- `*_EXAMPLE.md` +- `*_SUMMARY.md` +- Any other .md files without explicit user request except README.md files for new modules. + +✅ **Do document**: +- In-code docstrings +- README files for new modules +- Inline comments for complex logic +- Update existing README.md files when needed +- API documentation in existing docs structure + +## Backward Compatibility + +**Backward compatibility is NOT a priority** for this project: + +- Feel free to make breaking changes to improve code quality +- Remove deprecated code without extensive migration paths +- Update APIs for clarity and consistency +- Prioritize modern, clean design over legacy support + +## Code Style + +- Use descriptive variable and function names +- Add type hints to Python functions +- Keep functions focused and small +- Use `logging` module instead of `print` statements +- Follow PEP 8 for Python code + +## Python Commands + +**Use `py` for running Python on this Windows system:** + +```bash +# Running Python +py script.py + +# Running modules +py -m pytest tests/ + +# Python version +py --version +``` + +❌ **Don't use**: `python` (not available in PATH) +✅ **Do use**: `py` (Python launcher for Windows) + +## Testing + +- Update existing tests when changing APIs +- Use meaningful test names that describe what is being tested + +## Git Workflow + +**Do NOT stage files automatically:** + +❌ **Don't use**: `git add`, `git stage` +✅ **Do use**: `git status` to show what changed +✅ **User will**: Stage files themselves when ready + +The user prefers to review and stage changes manually. + +**Code and documentation versions:** +- Refer to code and documentation in the folder reference_code to get examples and documentation of the APIs and best practices for the advanced libraries used in this project: + - ITK, VTK, PyVista, Omniverse, PhysicsNeMo, Simpleware, MONAI, and OpenUSD diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6a3190e..f231667 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -382,17 +382,19 @@ jobs: - name: Install dev dependencies run: | python -m pip install --upgrade pip - pip install ruff mypy + # Ruff runs via pre-commit so its version is pinned in one place + # (.pre-commit-config.yaml); mypy is still invoked directly below. + pip install pre-commit mypy pip install -e ".[dev]" - - name: Check formatting with Ruff + - name: Check formatting with Ruff (via pre-commit) run: | - ruff format --check . + pre-commit run ruff-format --all-files continue-on-error: false - - name: Lint with Ruff + - name: Lint with Ruff (via pre-commit) run: | - ruff check . + pre-commit run ruff-check --all-files continue-on-error: false - name: Type check with mypy diff --git a/.gitignore b/.gitignore index 3e12bf9..c94f46d 100644 --- a/.gitignore +++ b/.gitignore @@ -13,7 +13,6 @@ build __pycache__ _version.py htmlcov -.cursor .vscode/* !.vscode/settings.json diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c415b09..95a2b3c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ repos: # Ruff - fast linter and formatter (replaces black, isort, flake8) - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.14 + rev: v0.15.20 hooks: # Run the linter with auto-fixes - id: ruff-check diff --git a/CLAUDE.md b/CLAUDE.md index 6eeb75a..b3a3f78 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,21 +2,82 @@ Project guidance for Claude Code in this repository. -## Role: -We are developing open-source code for scientific AI libraries. Leverage GPU-accelerated methods when appropriate. - -## Priorities (Ordered) -1) accuracy -2) clarity/maintainability/simplicity -3) consistency with the rest of the platform and open source standards -4) documentation -5) testing - -## Behavior Guidelines -1) Don't assume. Don't hide confusion. Surface tradeoffs. -2) Minimum code that solves the problem. Nothing speculative. -3) Touch only what you must. Clean up only your own mess. -4) Define success criteria. Loop until verified. +## Role + +We are developing open-source code for scientific AI libraries. Leverage +GPU-accelerated methods when appropriate. + +## 1. Think Before Coding + +**Don't assume. Don't hide confusion. Surface tradeoffs.** + +Before implementing: + +- State your assumptions explicitly. If uncertain, ask. +- If multiple interpretations exist, present them - don't pick silently. +- If a simpler approach exists, say so. Push back when warranted. +- If something is unclear, stop. Name what's confusing. Ask. + +## 2. Simplicity First + +**Minimum code that solves the problem. Nothing speculative.** + +- No features beyond what was asked. +- No abstractions for single-use code. +- No "flexibility" or "configurability" that wasn't requested. +- No error handling for impossible scenarios. +- If you write 200 lines and it could be 50, rewrite it. + +Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify. + +## 3. Surgical Changes + +**Touch only what you must. Clean up only your own mess.** + +When editing existing code: + +- Don't "improve" adjacent code, comments, or formatting. +- Don't refactor things that aren't broken. +- Match existing style, even if you'd do it differently. +- If you notice unrelated dead code, mention it - don't delete it. + +When your changes create orphans: + +- Remove imports/variables/functions that YOUR changes made unused. +- Don't remove pre-existing dead code unless asked. + +The test: Every changed line should trace directly to the user's request. + +## 4. Goal-Driven Execution + +**Define success criteria. Loop until verified.** + +Transform tasks into verifiable goals: + +- "Add validation" → "Write tests for invalid inputs, then make them pass" +- "Fix the bug" → "Write a test that reproduces it, then make it pass" +- "Refactor X" → "Ensure tests pass before and after" + +For multi-step tasks, state a brief plan: + +``` +1. [Step] → verify: [check] +2. [Step] → verify: [check] +3. [Step] → verify: [check] +``` + +Strong success criteria let you loop independently. Weak criteria +("make it work") require constant clarification. + +## 5. Project-Specific Rules + +- Breaking changes are acceptable. Backward-compatibility shims are not. + +--- + +**These guidelines are working if:** fewer unnecessary changes in diffs, +fewer rewrites due to overcomplication, and clarifying questions come +before implementation rather than after mistakes. ## Commands @@ -24,7 +85,7 @@ We are developing open-source code for scientific AI libraries. Leverage GPU-acc ```bash # Install in editable mode (preferred) -uv pip install -e . +uv pip install -e .[dev] # Lint and format ruff check . --fix && ruff format . @@ -35,36 +96,23 @@ mypy src/ tests/ # All pre-commit hooks pre-commit run --all-files -# Fast tests (recommended for development — slow/GPU/Simpleware/experiment -# /tutorial tests are auto-skipped unless their opt-in flag is passed) -py -m pytest tests/ -v - -# Single test file or test by name -py -m pytest tests/test_contour_tools.py -v -py -m pytest tests/test_contour_tools.py::test_extract_surface -v - -# With coverage -py -m pytest tests/ --cov=src/physiotwin4d --cov-report=html - -# Create missing baselines -py -m pytest tests/ --create-baselines -``` - **Version bumping:** `bumpver update --patch` (or `--minor`, `--major`) ## Architecture -All classes inherit from `PhysioTwin4DBase` (`physiotwin4d_base.py`), which provides -a shared logger. Use `self.log_info()`, `self.log_debug()` — never `print()`. +All classes inherit from `PhysioTwin4DBase` (`physiotwin4d_base.py`), +which provides a shared logger. Use `self.log_info()`, `self.log_debug()` +— never `print()`. -Consult `docs/API_MAP.md` for the full index of classes, methods, and signatures. -Regenerate it after any public API change: `py utils/generate_api_map.py` +Consult `docs/API_MAP.md` and graphify (see section below) for the full +index of classes, methods, and signatures. Regenerate API_MAP.md after any public API change: `py utils/generate_api_map.py` **Key data conventions:** + - Images: `itk.Image`, axes X, Y, Z [, T] in LPS world space (ITK's native frame; `itk.imread` normalizes DICOM, NIfTI, MHA, and NRRD inputs to LPS) stored using itk.imwrite with compression=True -- 4D time series: shape `(X, Y, Z, T)` — never silently squeeze or permute axes +— never silently squeeze or permute axes - Surfaces: `pv.PolyData` in LPS (inherited from the source `itk.Image` via `itk.vtk_image_from_image`); converted to USD right-handed Y-up only at USD export by `vtk_to_usd.lps_points_to_usd` (USD +X=Left, +Y=Superior, +Z=Anterior) @@ -74,6 +122,9 @@ Regenerate it after any public API change: `py utils/generate_api_map.py` ## Testing +- Fast tests (recommended for development — slow/GPU/Simpleware/experiment + /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.) @@ -83,21 +134,6 @@ Regenerate it after any public API change: `py utils/generate_api_map.py` - Prefer images from `ROOT/data/test/slicer_heart_small` for tests - Prefer storing results in subdirs `./results/` -## Working Process - -Before editing any code: -1. Read the relevant source file(s) in full. -2. Summarize current behavior in 2–4 sentences. -3. Identify success criteria / metrics -4. Refer to *_tools.py files for commonly used routines -5. Refer to workspace/reference_code (when available) for third-party libraries -6. Propose a numbered plan; confirm before implementing non-trivial changes. -7. Follow the behavior guidelines given above. -8. Implement in small, reviewable diffs. -9. Update docstrings and tests for every changed public method. -10. Do not commit changes or make pull requests unless specifically told to do so. - -Breaking changes are acceptable. Backward-compatibility shims are not. ## Agents and Skills @@ -107,7 +143,8 @@ Claude, Codex, and other AI tooling. - `/plan` — inspect files, summarize design, produce a numbered plan (no code changes) - `/impl` — read → summarize → plan → implement in small diffs -- `/test-feature` — propose test plan, write real-data-driven pytest tests with baselines +- `/test-feature` — propose test plan, write real-data-driven pytest tests + with baselines - `/doc-feature` — update docstrings (and remind you to run `/regen-api-map`) - `/regen-api-map` — regenerate `docs/API_MAP.md` and report public-API changes - `/check-conventions` — audit changed files against project hard rules @@ -135,3 +172,14 @@ Document via docstrings and inline comments. - Breaking changes are acceptable — backward compatibility is not a priority - Max line length: 88 characters - Follow behavior guidelines. + +## graphify + +This project has a knowledge graph at graphify-out/ with god nodes, community structure, and cross-file relationships. + +Rules: + +- For codebase questions, first run `graphify query ""` when graphify-out/graph.json exists. Use `graphify path "" ""` for relationships and `graphify explain ""` for focused concepts. These return a scoped subgraph, usually much smaller than GRAPH_REPORT.md or raw grep output. +- If graphify-out/wiki/index.md exists, use it for broad navigation instead of raw source browsing. +- Read graphify-out/GRAPH_REPORT.md only for broad architecture review or when query/path/explain do not surface enough context. +- After modifying code, run `graphify update .` to keep the graph current (AST-only, no API cost). diff --git a/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py b/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py new file mode 100644 index 0000000..ac1ec8f --- /dev/null +++ b/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py @@ -0,0 +1,203 @@ +"""Apply the trained PhysicsNeMo MeshGraphNet cardiac motion model to Case1Pack. + +Purpose +------- +Drive the cardiac MeshGraphNet (MGN) motion model trained by +``tutorial_09_byod_train_physicsnemo_mgn.py`` over the DirLab ``Case1Pack`` +heart and, for each requested cardiac stage, rasterize the inferred per-vertex +displacement onto the ``Case1Pack`` image grid as a deformation field via +:meth:`WorkflowInferPhysicsNeMoMGN.create_deformation_field`. + +Inputs +------ +- MGN model: the epoch-300 checkpoint in ``output/tutorial_09_byod_mgn_3``. + That run directory holds only intermittent epoch checkpoints (training did + not finalize), so the finalize-time inference assets are assembled here from + the self-describing epoch checkpoint plus the source PCA template. +- Reference image (defines the deformation-field grid): + ``data/DirLab-4DCT/Case1Pack_T70.mha``. +- Patient heart shape: the 15-mode ``pca-vol-kcl`` fit produced by re-running + Tutorial 5 against Case1Pack: + * coefficients: ``output/tutorial_05_heart_to_lung/ + tutorial_05_heart_to_lung_registered_coefficients.json`` + * patient-space mesh (binning positions): ``output/tutorial_05_heart_to_lung/ + tutorial_05_heart_to_lung_template_mesh_registered.vtu`` + +Coordinate-frame note +--------------------- +The MGN was trained in the ``pca-vol-kcl`` basis, so its PCA reconstruction of a +subject lands in the model's *canonical* frame (near the origin). The Tutorial 5 +fit aligned that model to Case1Pack with a pose transform that the shape +coefficients do not carry, so the patient heart sits in the scanner frame of +``Case1Pack_T70``. The displacements the network predicts depend only on the +coefficients and the stage, not on where the surface sits; therefore this script +passes the patient-space registered *mesh* as ``reference_surface`` so the +displacements are binned at positions that actually overlap the reference image. +The registered volume mesh is passed (not the surface ``.vtp``) because its +extracted surface reproduces the exact mean-shape point ordering the network +outputs, keeping the displacement/position correspondence intact. + +Caveat: this assumes the network's displacement vectors and the Case1Pack image +share a common anatomical (LPS) orientation, which holds for supine cardiac/lung +CT but is not re-derived here. + +Outputs (under ``output/temp_beating_heart``) +-------------------------------------------- +Per stage ``s`` (percent of the RR interval): +- ``deformation_field_s.mha`` - ITK vector image of ``(dx, dy, dz)`` mm +- ``surface_normal_field_s.mha`` - ITK vector image of reference normals +- ``deformed_heart_surface_s.vtp`` - the patient-space heart surface at + that stage (reference surface displaced by the network) + +The per-stage surfaces are then assembled, in stage order, into a single +animated 4D ``beating_heart.usd`` with a heart anatomy material via +:class:`WorkflowConvertVTKToUSD`. +""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path + +import itk +import numpy as np +import pyvista as pv + +from physiotwin4d import 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). The ``tutorial_09_byod_mgn_3`` directory holds only epoch + checkpoints, so this regenerates 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 = 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)) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + # MGN model produced by tutorial_09_byod_train_physicsnemo_mgn.py. + model_dir = tutorials_dir / "output" / "tutorial_09_byod_mgn_3" + epoch = 300 + # Source PCA template used to train the model (regenerates missing assets). + pca_mean_volume = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + + # Case1Pack reference image (defines the deformation-field grid). + reference_image_file = repo_root / "data" / "DirLab-4DCT" / "Case1Pack_T70.mha" + + # 15-mode pca-vol-kcl fit of the Case1Pack heart (Tutorial 5, re-run). + tutorial_05_dir = tutorials_dir / "output" / "tutorial_05_heart_to_lung" + coefficients_file = ( + tutorial_05_dir / "tutorial_05_heart_to_lung_registered_coefficients.json" + ) + registered_mesh_file = ( + tutorial_05_dir / "tutorial_05_heart_to_lung_template_mesh_registered.vtu" + ) + + output_dir = tutorials_dir / "output" / "temp_beating_heart" + output_dir.mkdir(parents=True, exist_ok=True) + + # Cardiac stages (fraction of the RR interval) to render as a beating heart. + stages = [round(0.1 * k, 2) for k in range(10)] + + for required in ( + reference_image_file, + coefficients_file, + registered_mesh_file, + pca_mean_volume, + ): + if not required.exists(): + raise FileNotFoundError(f"Required input not found: {required}") + + _ensure_mgn_inference_assets(model_dir, epoch, pca_mean_volume) + + reference_image = itk.imread(str(reference_image_file)) + + infer = WorkflowInferPhysicsNeMoMGN(model_directory=model_dir, epoch=epoch) + + deformation_files: list[Path] = [] + normal_files: list[Path] = [] + surface_files: list[Path] = [] + deformed_surfaces: list[pv.PolyData] = [] + for stage in stages: + result = infer.create_deformation_field( + shape_parameters=coefficients_file, + stage=float(stage), + reference_image=reference_image, + reference_surface=registered_mesh_file, + ) + pct = int(round(stage * 100)) + field_path = output_dir / f"deformation_field_s{pct:03d}.mha" + normal_path = output_dir / f"surface_normal_field_s{pct:03d}.mha" + surface_path = output_dir / f"deformed_heart_surface_s{pct:03d}.vtp" + itk.imwrite(result["deformation_field"], str(field_path), compression=True) + itk.imwrite(result["normal_image"], str(normal_path), compression=True) + result["deformed_surface"].save(str(surface_path)) + deformation_files.append(field_path) + normal_files.append(normal_path) + surface_files.append(surface_path) + deformed_surfaces.append(result["deformed_surface"]) + logging.info("stage %.2f -> %s", stage, field_path.name) + + # Assemble the ordered per-stage surfaces into one animated 4D USD. The + # stages loop over one RR interval, so play them back at one beat per second. + usd_workflow = WorkflowConvertVTKToUSD( + input_meshes=deformed_surfaces, + usd_project_name="beating_heart", + output_directory=output_dir, + appearance="anatomy", + anatomy_type="heart", + separate_by_connectivity=True, + frames_per_second=float(len(stages)), + log_level=logging.INFO, + ) + usd_result = usd_workflow.process() + + tutorial_results = { + "deformation_fields": deformation_files, + "normal_fields": normal_files, + "deformed_surfaces": surface_files, + "usd_file": usd_result["usd_file"], + } diff --git a/experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py b/experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py new file mode 100644 index 0000000..6c3c572 --- /dev/null +++ b/experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py @@ -0,0 +1,263 @@ +"""Combine respiratory and cardiac motion into a 4D surface sequence and USD. + +Purpose +------- +Produce a sequence of combined-motion surfaces for the DirLab ``Case1Pack`` +thorax by composing two motions on one static patient surface: + +- **Cardiac** motion from the ``temp_heart_and_lungs_beating_heart`` deformation + fields (one per cardiac stage), which displace the heart region. These fields + are defined on the reference (``T70``) grid. +- **Respiratory** motion from the ``tutorial_01_lung`` forward transforms + (one per breath phase), which warp the reference-frame surface to each phase + of the breathing cycle. + +One cardiac cycle (10 stages) is rendered per breath phase, for +``10 phases x 10 stages = 100`` output frames. + +Composition order +----------------- +The cardiac deformation is applied **first, at the reference frame** where the +cardiac fields are defined, and the breathing transform is then applied to the +already cardiac-deformed surface:: + + p_combined = forward_r( p + cardiac_s(p) ) + +This carries the heart's cardiac deformation into each respiratory phase (the +breathing deformation is applied to the heart deformation), rather than sampling +the reference-frame cardiac field at breathing-displaced positions where it is +not valid. + +Smooth respiratory progression +------------------------------ +Rather than holding the breath fixed for a whole cardiac cycle, the respiratory +transform is interpolated from the current breath phase toward the next by the +fraction of the way through the cardiac cycle (``stage / n_stages``). The +interpolation is a per-vertex linear blend of the two phases' warped positions, +equivalent to interpolating the two displacement transforms:: + + p_combined = (1 - t) * forward_r(cs) + t * forward_{r+1}(cs), t = stage / n_stages + +with ``cs = p + cardiac_s(p)``. The next phase wraps around (phase 9 -> phase 0) +so the breathing loops smoothly. + +Pipeline +-------- +1. Smooth each of the 10 cardiac deformation fields + (``output/temp_beating_heart/deformation_field_s0*.mha``) with a Gaussian of + ``SMOOTHING_SIGMA_MM`` (10 mm) into a ``DisplacementFieldTransform``. The raw + fields are a thin surface shell, so smoothing spreads them into a continuous + deformation (and, as a side effect, attenuates the peak magnitude). +2. Load the combined patient surface + (``output/tutorial_02_lung/patient_surfaces.vtp``), optionally decimate and + smooth it once (``SURFACE_DECIMATION_REDUCTION`` / + ``SURFACE_SMOOTHING_ITERATIONS``; both disabled by default), then + cardiac-deform it once per stage (reused across all breath phases). +3. Warp each cardiac-deformed surface to every breath phase with the + ``tutorial_01_lung`` forward transforms ``slice_{r:03d}_all_forward.hdf``. +4. Blend consecutive breath phases per frame and assemble the 100 frames into a + single animated 4D USD. + +Outputs (under ``output/temp_heart_and_lungs_combined``) +------------------------------------------------------- +- ``combined_frame_.vtp`` for ``iii`` in ``000 .. 099`` - the combined + respiratory + cardiac surface, ordered breath-phase-major then cardiac. +- ``heart_and_lungs_combined.usd`` - the 100-frame animated 4D USD. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import itk +import numpy as np +import pyvista as pv + +from physiotwin4d import ImageTools, TransformTools, WorkflowConvertVTKToUSD + +# Gaussian sigma (mm) used to smooth the sparse cardiac deformation fields. +SMOOTHING_SIGMA_MM = 10.0 +# USD playback rate; 10 stages/second ~= one heartbeat per second. +FRAMES_PER_SECOND = 10.0 +# One-time conditioning of the patient surface, applied before any warping. +# Fraction of triangles to remove (0.0 = no decimation, e.g. 0.5 halves them). +SURFACE_DECIMATION_REDUCTION = 0.0 +# Taubin (non-shrinking) smoothing iterations (0 = no smoothing). +SURFACE_SMOOTHING_ITERATIONS = 0 + +_transform_tools = TransformTools() + + +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 patient surface so every warped frame inherits the same + resolution and smoothing. Decimation uses ``decimate_pro`` on a triangulated + copy; smoothing uses non-shrinking Taubin smoothing. + """ + conditioned = surface + if decimation_reduction > 0.0: + conditioned = conditioned.triangulate().decimate_pro(decimation_reduction) + if smoothing_iterations > 0: + conditioned = conditioned.smooth_taubin(n_iter=smoothing_iterations) + return conditioned + + +def _smoothed_cardiac_transform( + field_file: Path, sigma_mm: float +) -> itk.DisplacementFieldTransform: + """Load a cardiac deformation field and return a smoothed field transform. + + The ``.mha`` field is a float vector image; it is converted to a + double-precision vector field, wrapped as a ``DisplacementFieldTransform``, + and Gaussian-smoothed by ``sigma_mm`` (in physical millimeters). + """ + field = itk.imread(str(field_file)) + 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 + ) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + logger = logging.getLogger("temp_heart_and_lungs_combined") + + tutorials_dir = Path(__file__).resolve().parent + + beating_heart_dir = tutorials_dir / "output" / "temp_beating_heart" + respiratory_dir = tutorials_dir / "output" / "tutorial_01_lung" + patient_surface_file = ( + tutorials_dir / "output" / "tutorial_02_lung" / "patient_surfaces.vtp" + ) + output_dir = tutorials_dir / "output" / "temp_heart_and_lungs_combined" + output_dir.mkdir(parents=True, exist_ok=True) + for stale_frame in output_dir.glob("combined_frame_*.vtp"): + stale_frame.unlink() + + cardiac_field_files = sorted(beating_heart_dir.glob("deformation_field_s*.mha")) + forward_transform_files = sorted(respiratory_dir.glob("slice_*_all_forward.hdf")) + + if not cardiac_field_files: + raise FileNotFoundError( + f"No cardiac deformation fields found in {beating_heart_dir}. " + "Run temp_heart_and_lungs_beating_heart.py first." + ) + if not forward_transform_files: + raise FileNotFoundError( + f"No respiratory forward transforms found in {respiratory_dir}. " + "Run tutorial_01_lung_gated_ct_to_usd.py first." + ) + if not patient_surface_file.exists(): + raise FileNotFoundError( + f"Patient surface not found: {patient_surface_file}. " + "Run tutorial_02_lung_ct_to_vtk.py first." + ) + + n_phases = len(forward_transform_files) + n_stages = len(cardiac_field_files) + logger.info( + "Respiratory phases: %d, cardiac stages: %d (1 cardiac cycle/phase)", + n_phases, + n_stages, + ) + + # Smooth every cardiac field once; reused across all breath phases. + cardiac_transforms = [ + _smoothed_cardiac_transform(field_file, SMOOTHING_SIGMA_MM) + for field_file in cardiac_field_files + ] + logger.info( + "Smoothed %d cardiac deformation fields by %.1f mm", + len(cardiac_transforms), + SMOOTHING_SIGMA_MM, + ) + + patient_surface = pv.read(str(patient_surface_file)) + # Condition the model once (decimate/smooth) so every frame inherits it. + patient_surface = _condition_surface( + patient_surface, SURFACE_DECIMATION_REDUCTION, SURFACE_SMOOTHING_ITERATIONS + ) + logger.info( + "Patient surface: %d points (decimation=%.2f, smoothing_iters=%d)", + patient_surface.n_points, + SURFACE_DECIMATION_REDUCTION, + SURFACE_SMOOTHING_ITERATIONS, + ) + + # Cardiac motion is applied at the reference frame (where the fields are + # defined), once per stage, then carried into each breath phase below. + cardiac_surfaces = [ + _transform_tools.transform_pvcontour(patient_surface, cardiac_transform) + for cardiac_transform in cardiac_transforms + ] + + # 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)) + resp_points.append( + [ + np.asarray( + _transform_tools.transform_pvcontour( + cardiac_surface, forward_transform + ).points, + dtype=np.float32, + ) + for cardiac_surface in cardiac_surfaces + ] + ) + logger.info("Respiratory warp phase %d/%d done", phase_idx + 1, n_phases) + + # Blend the current and next breath phase by the fraction through the cardiac + # cycle, so breathing advances smoothly across the heartbeat (next wraps). + combined_files: list[Path] = [] + usd_frames: list[pv.PolyData] = [] + for phase_idx in range(n_phases): + next_phase_idx = (phase_idx + 1) % n_phases + for stage_idx in range(n_stages): + blend = stage_idx / n_stages + points = (1.0 - blend) * resp_points[phase_idx][stage_idx] + ( + blend * resp_points[next_phase_idx][stage_idx] + ) + combined_surface = patient_surface.copy(deep=True) + combined_surface.points = points + + frame_idx = phase_idx * n_stages + stage_idx + frame_file = output_dir / f"combined_frame_{frame_idx:03d}.vtp" + combined_surface.save(str(frame_file)) + combined_files.append(frame_file) + usd_frames.append(combined_surface) + + del resp_points + logger.info( + "Wrote %d combined-motion surfaces to %s", len(combined_files), output_dir + ) + + # Assemble the ordered frames into a single animated 4D USD. + usd_workflow = WorkflowConvertVTKToUSD( + input_meshes=usd_frames, + usd_project_name="heart_and_lungs_combined", + output_directory=output_dir, + appearance="solid", + solid_color=(0.82, 0.70, 0.66), + separate_by_connectivity=False, + frames_per_second=FRAMES_PER_SECOND, + ) + usd_result = usd_workflow.process() + logger.info("Wrote 4D USD: %s", usd_result["usd_file"]) + + tutorial_results = { + "combined_surfaces": combined_files, + "usd_file": usd_result["usd_file"], + } diff --git a/src/physiotwin4d/segment_chest_total_segmentator.py b/src/physiotwin4d/segment_chest_total_segmentator.py index 098d79d..e317f7f 100644 --- a/src/physiotwin4d/segment_chest_total_segmentator.py +++ b/src/physiotwin4d/segment_chest_total_segmentator.py @@ -179,6 +179,7 @@ def __init__(self, log_level: int | str = logging.INFO): 115: "rib_right_12", 116: "sternum", 117: "costal_cartilages", + 25: "sacrum", }, ), ( @@ -199,7 +200,6 @@ def __init__(self, log_level: int | str = logging.INFO): 20: "colon", 21: "urinary_bladder", 22: "prostate", - 25: "sacrum", 80: "gluteus_maximus_left", 81: "gluteus_maximus_right", 82: "gluteus_medius_left", @@ -380,8 +380,8 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: # labelmap_arr_body contains: 1=body, 2=body_trunc, 3=body_extremities, # 4=skin # Only overwrite the background with body labels - mask = final_arr > 0 - labelmap_arr_body[mask] = 0 + # mask = final_arr > 0 + # labelmap_arr_body[mask] = 0 final_arr = np.where( labelmap_arr_body == 4, 133, final_arr ) # body_skin diff --git a/src/physiotwin4d/usd_anatomy_tools.py b/src/physiotwin4d/usd_anatomy_tools.py index 0efd5d8..a6f6de7 100644 --- a/src/physiotwin4d/usd_anatomy_tools.py +++ b/src/physiotwin4d/usd_anatomy_tools.py @@ -30,7 +30,7 @@ """ import logging -from typing import Any, Mapping +from typing import Any, Mapping, Optional from pxr import Sdf, UsdGeom, UsdShade @@ -38,10 +38,12 @@ # Default OmniSurface render parameters keyed by group name (matching # :class:`physiotwin4d.AnatomyTaxonomy.group_names`) and by organ-level -# overrides (e.g. ``liver``, ``spleen``, ``kidney_left``). ``enhance_meshes`` -# consults an organ-level entry first, then falls back to the containing -# group's entry, and finally to ``"other"``. Module-level so CLIs and tests -# can enumerate the supported types without constructing a USD stage. +# overrides (e.g. ``liver``, ``spleen``, ``kidney``). ``enhance_meshes`` +# consults the organ-level entries first (matching an override key when it is +# a substring of the organ name, so ``kidney`` covers ``kidney_left`` and +# ``kidney_right``), then falls back to the containing group's entry, and +# finally to ``"other"``. Module-level so CLIs and tests can enumerate the +# supported types without constructing a USD stage. DEFAULT_RENDER_PARAMS: dict[str, dict[str, Any]] = { "heart": { "name": "Heart", @@ -148,10 +150,14 @@ "subsurface_scale": 0.3, "coat_weight": 0.12, }, - # Organ-level overrides. enhance_meshes consults these by organ - # name *before* falling back to the containing group's params, so - # liver/spleen/kidney get their dedicated look despite being in - # the soft_tissue group of the taxonomy. + # Organ-level overrides. enhance_meshes consults these before falling + # back to the containing group's params, matching an override key when it + # is a substring of the organ name, so e.g. liver/spleen/kidney get their + # dedicated look despite being in the soft_tissue group of the taxonomy. + # The overrides below span abdominal/endocrine organs, vessels (artery vs. + # vein), tissue types (skin/fat/muscle/cartilage), and the oxygenation- + # coded heart chambers; when several keys are substrings of one name the + # longest (most specific) one wins. "liver": { "name": "Liver", "enable_diffuse_transmission": True, @@ -182,7 +188,7 @@ "subsurface_scale": 0.15, "coat_weight": 0.1, }, - "kidney_right": { + "kidney": { "name": "Kidney", "enable_diffuse_transmission": True, "diffuse_reflection_weight": 0.1, @@ -197,11 +203,421 @@ "subsurface_scale": 0.18, "coat_weight": 0.1, }, + # Skin: tan/pink dermis. Skin is the canonical strong-subsurface tissue + # (Jensen et al., "A Practical Model for Subsurface Light Transport", + # SIGGRAPH 2001), so it carries the highest subsurface_weight here plus a + # thin oily coat for the epidermal sheen. Matches organs named "*skin*". + "skin": { + "name": "Skin", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.2, + "diffuse_reflection_color": (0.75, 0.52, 0.42), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.45, + "subsurface_transmission_color": (0.85, 0.5, 0.4), + "subsurface_scattering_color": (0.9, 0.55, 0.45), + "subsurface_weight": 0.12, + "subsurface_scale": 0.4, + "coat_weight": 0.15, + }, + # Airway: pale-pink tracheobronchial mucosa. Wet mucus surface -> strong, + # low-roughness specular plus a heavy coat for the glossy sheen. Matches + # organs named "*airway*" (e.g. lung_airways, lung_airways_wall). + "airway": { + "name": "Airway", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.18, + "diffuse_reflection_color": (0.85, 0.6, 0.58), + "diffuse_reflection_roughness": 0.4, + "metalness": 0.0, + "specular_reflection_weight": 0.6, + "specular_reflection_roughness": 0.35, + "subsurface_transmission_color": (0.9, 0.7, 0.7), + "subsurface_scattering_color": (0.9, 0.7, 0.7), + "subsurface_weight": 0.08, + "subsurface_scale": 0.15, + "coat_weight": 0.2, + }, + # Vein: deoxygenated blood -> dark, desaturated purplish-red (blue channel + # lifted above green). Deoxy-hemoglobin absorbs green/red more than oxy-Hb, + # shifting venous vessels darker and bluer than arteries. Matches organs + # named "*vein*" (pulmonary_vein, brachiocephalic_vein, lung_veins). + "vein": { + "name": "Vein", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.2, + "diffuse_reflection_color": (0.3, 0.03, 0.12), + "diffuse_reflection_roughness": 0.4, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.4, + "subsurface_transmission_color": (0.5, 0.1, 0.2), + "subsurface_scattering_color": (0.5, 0.1, 0.2), + "subsurface_weight": 0.09, + "subsurface_scale": 0.2, + "coat_weight": 0.12, + }, + # Artery: oxygenated blood -> brighter, more saturated red than veins. Key + # is the substring "arter" so it matches both "artery" and "arteries" + # (subclavian/carotid arteries, lung_arteries, pulmonary_artery). + "arter": { + "name": "Artery", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.55, 0.03, 0.03), + "diffuse_reflection_roughness": 0.35, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.4, + "subsurface_transmission_color": (0.8, 0.2, 0.2), + "subsurface_scattering_color": (0.8, 0.2, 0.2), + "subsurface_weight": 0.09, + "subsurface_scale": 0.22, + "coat_weight": 0.12, + }, + # Fat: adipose tissue -> pale yellow, translucent, with a greasy coat. + # Matches organs named "*fat*" (subcutaneous, torso, intermuscular). + "fat": { + "name": "Fat", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.2, + "diffuse_reflection_color": (0.9, 0.75, 0.35), + "diffuse_reflection_roughness": 0.45, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.4, + "subsurface_transmission_color": (0.95, 0.85, 0.5), + "subsurface_scattering_color": (0.95, 0.85, 0.5), + "subsurface_weight": 0.1, + "subsurface_scale": 0.3, + "coat_weight": 0.15, + }, + # Muscle: skeletal muscle -> deep red-brown, matte (fibrous striations), + # only lightly moist. Matches organs named "*muscle*" (skeletal_muscle); + # the substring does not appear in "intermuscular_fat", so fat is + # unaffected. Darker and less saturated than arterial blood. + "muscle": { + "name": "Muscle", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.18, + "diffuse_reflection_color": (0.42, 0.08, 0.07), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.6, 0.15, 0.13), + "subsurface_scattering_color": (0.6, 0.15, 0.13), + "subsurface_weight": 0.08, + "subsurface_scale": 0.18, + "coat_weight": 0.05, + }, + # Aorta: largest systemic artery, oxygenated -> bright saturated red, same + # optical look as the generic "arter" override but kept as its own entry + # because "aorta" contains no "arter" substring. Matches "aorta", + # "highres_aorta". + "aorta": { + "name": "Aorta", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.55, 0.03, 0.03), + "diffuse_reflection_roughness": 0.35, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.4, + "subsurface_transmission_color": (0.8, 0.2, 0.2), + "subsurface_scattering_color": (0.8, 0.2, 0.2), + "subsurface_weight": 0.09, + "subsurface_scale": 0.22, + "coat_weight": 0.12, + }, + # Vena cava: large systemic vein, deoxygenated -> dark purplish-red, same + # optical look as the generic "vein" override but kept as its own entry + # because "vena_cava" contains no "vein" substring. Matches + # "superior_vena_cava", "inferior_vena_cava". + "cava": { + "name": "Vena_Cava", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.2, + "diffuse_reflection_color": (0.3, 0.03, 0.12), + "diffuse_reflection_roughness": 0.4, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.4, + "subsurface_transmission_color": (0.5, 0.1, 0.2), + "subsurface_scattering_color": (0.5, 0.1, 0.2), + "subsurface_weight": 0.09, + "subsurface_scale": 0.2, + "coat_weight": 0.12, + }, + # --- Abdominal / endocrine organ overrides (soft_tissue group). These + # organs otherwise all collapse onto the generic tan soft_tissue look. + # Colors are realistic but nudged apart so neighboring organs stay + # distinguishable during cropped visual inspection. --- + # Gallbladder: bile-filled -> characteristic olive/green, the single most + # recognizable abdominal cue. Matches "gallbladder". + "gallbladder": { + "name": "Gallbladder", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.17, + "diffuse_reflection_color": (0.35, 0.45, 0.15), + "diffuse_reflection_roughness": 0.45, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.5, 0.6, 0.25), + "subsurface_scattering_color": (0.5, 0.6, 0.25), + "subsurface_weight": 0.09, + "subsurface_scale": 0.25, + "coat_weight": 0.15, + }, + # Pancreas: lobulated -> pale salmon-tan. Matches "pancreas". + "pancreas": { + "name": "Pancreas", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.17, + "diffuse_reflection_color": (0.8, 0.52, 0.44), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.9, 0.65, 0.55), + "subsurface_scattering_color": (0.9, 0.65, 0.55), + "subsurface_weight": 0.08, + "subsurface_scale": 0.25, + "coat_weight": 0.1, + }, + # Stomach: muscular wall -> pink-red, redder than the pale bowel/pancreas. + # Matches "stomach". + "stomach": { + "name": "Stomach", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.17, + "diffuse_reflection_color": (0.72, 0.4, 0.38), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.85, 0.5, 0.45), + "subsurface_scattering_color": (0.85, 0.5, 0.45), + "subsurface_weight": 0.08, + "subsurface_scale": 0.3, + "coat_weight": 0.12, + }, + # Brain: soft neural tissue -> light pink-grey with a little more + # subsurface glow than firm organs. Matches "brain". + "brain": { + "name": "Brain", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.17, + "diffuse_reflection_color": (0.82, 0.72, 0.68), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.9, 0.82, 0.8), + "subsurface_scattering_color": (0.9, 0.82, 0.8), + "subsurface_weight": 0.1, + "subsurface_scale": 0.3, + "coat_weight": 0.1, + }, + # Thyroid: highly vascular gland -> deep red-brown. Matches "thyroid". + "thyroid": { + "name": "Thyroid", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.15, + "diffuse_reflection_color": (0.55, 0.22, 0.2), + "diffuse_reflection_roughness": 0.45, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.7, 0.3, 0.25), + "subsurface_scattering_color": (0.7, 0.3, 0.25), + "subsurface_weight": 0.085, + "subsurface_scale": 0.2, + "coat_weight": 0.1, + }, + # Adrenal gland: cortex -> yellow-orange. More saturated/orange than the + # pale yellow "fat" look so the two stay distinct. Matches "adrenal". + "adrenal": { + "name": "Adrenal_Gland", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.17, + "diffuse_reflection_color": (0.82, 0.62, 0.32), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.9, 0.75, 0.5), + "subsurface_scattering_color": (0.9, 0.75, 0.5), + "subsurface_weight": 0.08, + "subsurface_scale": 0.25, + "coat_weight": 0.1, + }, + # Urinary bladder: thin, fluid-filled wall -> pale pink, lightly wet. + # Key "bladder" also occurs in "gallbladder", but override matching is + # longest-first, so "gallbladder" (above) wins for the gallbladder and + # "bladder" claims only "urinary_bladder". + "bladder": { + "name": "Urinary_Bladder", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.15, + "diffuse_reflection_color": (0.8, 0.62, 0.6), + "diffuse_reflection_roughness": 0.45, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.45, + "subsurface_transmission_color": (0.9, 0.75, 0.72), + "subsurface_scattering_color": (0.9, 0.75, 0.72), + "subsurface_weight": 0.09, + "subsurface_scale": 0.25, + "coat_weight": 0.15, + }, + # Gluteus muscles: skeletal muscle -> deep red-brown. Shares the "muscle" + # look but needs its own key because the gluteus label names contain no + # "muscle" substring. Matches "gluteus_maximus/medius/minimus_*". + "gluteus": { + "name": "Gluteus_Muscle", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.18, + "diffuse_reflection_color": (0.42, 0.08, 0.07), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.4, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.6, 0.15, 0.13), + "subsurface_scattering_color": (0.6, 0.15, 0.13), + "subsurface_weight": 0.08, + "subsurface_scale": 0.18, + "coat_weight": 0.05, + }, + # Costal cartilage: cartilage, not cortical bone -> translucent grey-white + # with more subsurface than the opaque "bone" look. Bone-group organ, but + # organ overrides win over the group. Matches "costal_cartilages". + "cartilage": { + "name": "Costal_Cartilage", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.2, + "diffuse_reflection_color": (0.82, 0.86, 0.82), + "diffuse_reflection_roughness": 0.35, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.85, 0.9, 0.85), + "subsurface_scattering_color": (0.85, 0.9, 0.85), + "subsurface_weight": 0.1, + "subsurface_scale": 0.2, + "coat_weight": 0.05, + }, + # --- Heart-chamber overrides (heart group). Oxygenation-coded so cardiac + # anatomy reads at a glance: left (systemic, oxygenated) = brighter red; + # right (pulmonary, deoxygenated) = darker, bluer red; myocardium = deep + # red-brown wall. Keys carry the full "_left"/"_right" suffix because the + # two sides are not substrings of one another. --- + "myocardium": { + "name": "Myocardium", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.4, 0.08, 0.07), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.6, 0.25, 0.22), + "subsurface_scattering_color": (0.6, 0.25, 0.22), + "subsurface_weight": 0.1, + "subsurface_scale": 1.5, + "coat_weight": 0.0, + }, + "atrium_left": { + "name": "Atrium_Left", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.5, 0.06, 0.06), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.85, 0.4, 0.4), + "subsurface_scattering_color": (0.85, 0.4, 0.4), + "subsurface_weight": 0.1, + "subsurface_scale": 2.0, + "coat_weight": 0.0, + }, + "ventricle_left": { + "name": "Ventricle_Left", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.45, 0.03, 0.03), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.8, 0.35, 0.35), + "subsurface_scattering_color": (0.8, 0.35, 0.35), + "subsurface_weight": 0.1, + "subsurface_scale": 2.0, + "coat_weight": 0.0, + }, + "atrium_right": { + "name": "Atrium_Right", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.32, 0.05, 0.12), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.6, 0.25, 0.35), + "subsurface_scattering_color": (0.6, 0.25, 0.35), + "subsurface_weight": 0.1, + "subsurface_scale": 2.0, + "coat_weight": 0.0, + }, + "ventricle_right": { + "name": "Ventricle_Right", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.28, 0.03, 0.11), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.55, 0.2, 0.32), + "subsurface_scattering_color": (0.55, 0.2, 0.32), + "subsurface_weight": 0.1, + "subsurface_scale": 2.0, + "coat_weight": 0.0, + }, + # Left atrial appendage: part of the left atrium -> oxygenated look. Key + # "atrial" is not a substring of "atrium", so it picks up only + # "atrial_appendage_left". + "atrial": { + "name": "Atrial_Appendage", + "enable_diffuse_transmission": True, + "diffuse_reflection_weight": 0.25, + "diffuse_reflection_color": (0.5, 0.06, 0.06), + "diffuse_reflection_roughness": 0.5, + "metalness": 0.0, + "specular_reflection_weight": 0.5, + "specular_reflection_roughness": 0.5, + "subsurface_transmission_color": (0.85, 0.4, 0.4), + "subsurface_scattering_color": (0.85, 0.4, 0.4), + "subsurface_weight": 0.1, + "subsurface_scale": 2.0, + "coat_weight": 0.0, + }, } -# kidney_left and kidney share the same look as kidney_right ("kidney" is a -# convenience alias kept for the CLI and other generic callers). -DEFAULT_RENDER_PARAMS["kidney_left"] = DEFAULT_RENDER_PARAMS["kidney_right"] -DEFAULT_RENDER_PARAMS["kidney"] = DEFAULT_RENDER_PARAMS["kidney_right"] + +# Canonical AnatomyTaxonomy group names that carry a group-level entry in +# DEFAULT_RENDER_PARAMS. Every other key is an organ-level override. Used by +# :meth:`USDAnatomyTools._resolve_render_params` to mirror ``enhance_meshes``, +# where an organ override always beats the containing group on a substring +# match (e.g. "lung_veins" -> the "vein" override, not the "lung" group). +GROUP_RENDER_KEYS: frozenset[str] = frozenset( + {"heart", "lung", "bone", "major_vessels", "contrast", "soft_tissue", "other"} +) class USDAnatomyTools(PhysioTwin4DBase): @@ -234,6 +650,34 @@ def get_anatomy_types(self) -> list[str]: """Return list of registered render-param keys (groups + organ overrides).""" return list(self.render_params.keys()) + def _resolve_render_params(self, anatomy_type: str) -> Optional[dict[str, Any]]: + """Resolve *anatomy_type* to a render-params entry, or ``None``. + + Matching mirrors :meth:`enhance_meshes`: an exact (case-insensitive) + key wins first; otherwise a key that is a substring of *anatomy_type* + matches (so ``"kidney_left"`` resolves to the ``"kidney"`` entry). + Organ-level overrides are tried before group-level keys (so + ``"lung_veins"`` resolves to the ``"vein"`` override, not the ``"lung"`` + group), and within each set the longest (most specific) key wins. + + Args: + anatomy_type: A group/organ name or registered render-params key. + + Returns: + The matching render-params dict, or ``None`` if nothing matches. + """ + name = anatomy_type.lower() + exact = self.render_params.get(name) + if exact is not None: + return exact + override_keys = [k for k in self.render_params if k not in GROUP_RENDER_KEYS] + group_keys = [k for k in self.render_params if k in GROUP_RENDER_KEYS] + for keys in (override_keys, group_keys): + for key in sorted(keys, key=len, reverse=True): + if key in name: + return self.render_params[key] + return None + def get_anatomy_diffuse_color( self, anatomy_type: str ) -> tuple[float, float, float]: @@ -243,16 +687,18 @@ def get_anatomy_diffuse_color( created with ``stage=None`` purely for color look-up purposes. Args: - anatomy_type: A registered render-params key (e.g. ``"heart"``, - ``"lung"``, ``"liver"``). + anatomy_type: A group/organ name or registered render-params key + (e.g. ``"heart"``, ``"lung"``, ``"liver"``). Resolved via + exact-then-substring matching, so ``"kidney_left"`` maps to the + ``"kidney"`` entry (see :meth:`_resolve_render_params`). Returns: RGB tuple of floats in ``[0, 1]``. Raises: - ValueError: If *anatomy_type* is not registered. + ValueError: If *anatomy_type* matches no registered entry. """ - params = self.render_params.get(anatomy_type.lower()) + params = self._resolve_render_params(anatomy_type) if params is None: raise ValueError( f"Unknown anatomy_type '{anatomy_type}'. " @@ -266,13 +712,15 @@ def apply_anatomy_material_to_mesh(self, mesh_path: str, anatomy_type: str) -> N Args: mesh_path: USD path to the mesh prim (e.g. "/World/Meshes/MyMesh"). - anatomy_type: A registered render-params key (group or organ - override). See :meth:`get_anatomy_types`. + anatomy_type: A group/organ name or registered render-params key. + Resolved via exact-then-substring matching, so ``"kidney_left"`` + maps to the ``"kidney"`` entry (see + :meth:`_resolve_render_params`). See :meth:`get_anatomy_types`. Raises: - ValueError: If mesh_path is invalid or anatomy_type is not registered. + ValueError: If mesh_path is invalid or anatomy_type matches no entry. """ - params = self.render_params.get(anatomy_type.lower()) + params = self._resolve_render_params(anatomy_type) if params is None: raise ValueError( f"Unknown anatomy_type '{anatomy_type}'. " @@ -362,8 +810,11 @@ def enhance_meshes(self, segmentator: Any) -> None: Walks the segmenter's :class:`AnatomyTaxonomy` and applies a material to each mesh prim whose leaf name matches an organ name in any group. An organ-level entry in :attr:`render_params` (e.g. ``"liver"``, - ``"spleen"``, ``"kidney_left"``) takes precedence over the entry for - the containing group. + ``"spleen"``, ``"kidney"``) takes precedence over the entry for the + containing group. An organ-level key matches when it is a substring of + the organ name, so ``"kidney"`` covers both ``kidney_left`` and + ``kidney_right``; when several keys match, the longest (most specific) + one wins. Anatomy grouping is performed upstream by ConvertVTKToUSD, which writes labeled prims under ``/World/{basename}/{type}/{label_name}``. @@ -375,18 +826,33 @@ def enhance_meshes(self, segmentator: Any) -> None: """ taxonomy = segmentator.taxonomy + # Organ-level override keys are every registered render-params key + # that is neither a group name nor the "other" fallback. They match an + # organ when the key is a substring of the organ name (e.g. "kidney" + # matches "kidney_left"). Sorted longest-first so a more specific key + # wins when several are substrings of the same organ name. + group_names = set(taxonomy.group_names()) + override_keys = sorted( + (k for k in self.render_params if k not in group_names and k != "other"), + key=len, + reverse=True, + ) + # Build organ_name -> render params dict in one pass. Organ-level - # overrides win over the containing group's params; if neither is - # registered, fall back to the "other" entry (always present in + # overrides win over the containing group's params; if neither + # applies, fall back to the "other" entry (always present in # DEFAULT_RENDER_PARAMS, so the lookup is safe). organ_params: dict[str, dict[str, Any]] = {} default_params: dict[str, Any] = self.render_params["other"] for group_name in taxonomy.group_names(): group_params = self.render_params.get(group_name, default_params) for organ_name in taxonomy.labels_in_group(group_name).values(): - organ_params[organ_name] = self.render_params.get( - organ_name, group_params - ) + selected = group_params + for key in override_keys: + if key in organ_name.lower(): + selected = self.render_params[key] + break + organ_params[organ_name] = selected for prim in self.stage.Traverse(): mesh_prim = UsdGeom.Mesh(prim) diff --git a/tutorials/tutorial_03_heart_vtk_to_usd.py b/tutorials/tutorial_03_heart_vtk_to_usd.py index 750b194..e2b9a34 100644 --- a/tutorials/tutorial_03_heart_vtk_to_usd.py +++ b/tutorials/tutorial_03_heart_vtk_to_usd.py @@ -44,9 +44,7 @@ project_name = "tutorial_02_heart" # Preferred input: the combined surface saved by Tutorial 2. - vtk_file = ( - tutorials_dir / "output" / "tutorial_02_heart" / "patient_surfaces.vtp" - ) + vtk_file = tutorials_dir / "output" / "tutorial_02_heart" / "patient_surfaces.vtp" log_level = logging.INFO @@ -86,4 +84,4 @@ ) ] - tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} \ No newline at end of file + tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} diff --git a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py index 83a5993..09eb083 100644 --- a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py @@ -103,19 +103,29 @@ segmentation_result = segmentation_method.segment(patient_image) patient_labelmap = segmentation_result["labelmap"] - itk.imwrite(patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz") + itk.imwrite( + patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz" + ) heart_labelmap = segmentation_result["heart"] - itk.imwrite(heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz") - + itk.imwrite( + heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz" + ) + contour_tools = ContourTools() heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) heart_surface.save(output_dir / f"{project_name}_heart_surface.vtp") else: - patient_labelmap = itk.imread(output_dir / f"{project_name}_patient_labelmap.nii.gz") - heart_labelmap = itk.imread(output_dir / f"{project_name}_heart_labelmap.nii.gz") - heart_surface = pv.read(output_dir / f"{project_name}_heart_surface.vtp") + patient_labelmap = itk.imread( + output_dir / f"{project_name}_patient_labelmap.nii.gz" + ) + heart_labelmap = itk.imread( + output_dir / f"{project_name}_heart_labelmap.nii.gz" + ) + heart_surface = cast( + pv.PolyData, pv.read(output_dir / f"{project_name}_heart_surface.vtp") + ) # Workflow initialization @@ -141,21 +151,31 @@ # Result saving registered_coefficients = workflow.pca_coefficients if registered_coefficients is not None: - registered_coefficients_path = output_dir / f"{project_name}_registered_coefficients.json" + registered_coefficients_path = ( + output_dir / f"{project_name}_registered_coefficients.json" + ) with registered_coefficients_path.open(mode="w", encoding="utf-8") as f: json.dump(registered_coefficients.tolist(), f) template_mesh = workflow.pca_template_model + assert template_mesh is not None, "pca_template_model must be set after process()" template_mesh.save(str(output_dir / f"{project_name}_template_mesh.vtu")) template_surface = workflow.pca_template_model_surface + assert template_surface is not None, ( + "pca_template_model_surface must be set after process()" + ) template_surface.save(str(output_dir / f"{project_name}_template_surface.vtp")) registered_mesh = workflow_results["registered_template_model"] - registered_mesh.save(str(output_dir / f"{project_name}_template_mesh_registered.vtu")) + registered_mesh.save( + str(output_dir / f"{project_name}_template_mesh_registered.vtu") + ) registered_surface = workflow_results["registered_template_model_surface"] - registered_surface.save(str(output_dir / f"{project_name}_template_surface_registered.vtp")) + registered_surface.save( + str(output_dir / f"{project_name}_template_surface_registered.vtp") + ) # Testing TestTools( diff --git a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py index a6276e5..8c2fc29 100644 --- a/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py @@ -90,7 +90,12 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" MANIFESTS_DIR = OUTPUT_DIR / "manifests_mgn" - RESUME_FROM = str(TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn_2" / "mgn_stage_model_epoch_00200.pt") + RESUME_FROM = str( + TUTORIALS_DIR + / "output" + / "tutorial_09_byod_mgn_2" + / "mgn_stage_model_epoch_00200.pt" + ) EPOCHS = 1500 BATCH_SIZE_GRAPHS = 4 # mini-batch measured in (subject, phase) graphs @@ -101,7 +106,7 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ # Explicit held-out splits; every other discovered subject is used for training. TEST_SUBJECTS = ["pm0027"] - VAL_SUBJECTS = [] + VAL_SUBJECTS: list[str] = [] LOG_LEVEL = logging.INFO """Discover subjects, train a MeshGraphNet, and evaluate the test split.""" @@ -155,10 +160,16 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ trainer.set_num_layers(NUM_LAYERS) train_result = trainer.process() - # Evaluate held-out test subjects against their ground-truth phases. - infer = WorkflowInferPhysicsNeMoMGN(model_directory=OUTPUT_DIR, log_level=LOG_LEVEL) + # Evaluate held-out test subjects against their ground-truth phases. When + # resuming, training writes to a fresh sibling directory, so evaluate the + # model from the directory training actually used, not the original + # OUTPUT_DIR. + model_directory = train_result["output_directory"] + infer = WorkflowInferPhysicsNeMoMGN( + model_directory=model_directory, log_level=LOG_LEVEL + ) eval_outputs: dict[str, Any] = {} for sid in TEST_SUBJECTS: eval_outputs[sid] = infer.predict( - manifests[sid], output_directory=OUTPUT_DIR / "eval_mgn" / sid - ) \ No newline at end of file + manifests[sid], output_directory=model_directory / "eval_mgn" / sid + ) diff --git a/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py b/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py index beaf89e..4f11200 100644 --- a/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py @@ -99,7 +99,7 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ # Explicit held-out splits; every other discovered subject is used for training. TEST_SUBJECTS = ["pm0027"] - VAL_SUBJECTS = [] + VAL_SUBJECTS: list[str] = [] LOG_LEVEL = logging.INFO def run_tutorial() -> dict[str, Any]: @@ -150,13 +150,16 @@ def run_tutorial() -> dict[str, Any]: trainer.set_num_layers(NUM_LAYERS) train_result = trainer.process() + # Evaluate from the directory training actually used: when resuming, + # training writes to a fresh sibling directory rather than OUTPUT_DIR. + model_directory = train_result["output_directory"] infer = WorkflowInferPhysicsNeMoMLP( - model_directory=OUTPUT_DIR, log_level=LOG_LEVEL + model_directory=model_directory, log_level=LOG_LEVEL ) eval_outputs: dict[str, Any] = {} for sid in TEST_SUBJECTS: eval_outputs[sid] = infer.predict( - manifests[sid], output_directory=OUTPUT_DIR / "eval_mlp" / sid + manifests[sid], output_directory=model_directory / "eval_mlp" / sid ) return {"training": train_result, "evaluation": eval_outputs} diff --git a/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py index eb65b1a..56c4b9e 100644 --- a/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py +++ b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py @@ -36,7 +36,7 @@ import logging import sys from pathlib import Path -from typing import Any, Optional, cast +from typing import Any, Optional from physiotwin4d import WorkflowInferPhysicsNeMoMGN @@ -92,10 +92,7 @@ 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 cast( - "dict[str, Any]", - infer.predict(manifest_path, stages=stages, output_directory=out_dir), - ) + return infer.predict(manifest_path, stages=stages, output_directory=out_dir) def run_tutorial() -> dict[str, Any]: diff --git a/tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py b/tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py index f7d1fb0..d3d7859 100644 --- a/tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py +++ b/tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py @@ -36,7 +36,7 @@ import logging import sys from pathlib import Path -from typing import Any, Optional, cast +from typing import Any, Optional from physiotwin4d import WorkflowInferPhysicsNeMoMLP @@ -92,10 +92,7 @@ def predict( """Predict cardiac surfaces for *subject* using the trained MLP.""" manifest_path = _write_subject_manifest(subject, out_dir) infer = WorkflowInferPhysicsNeMoMLP(model_directory=MODEL_DIR, epoch=epoch) - return cast( - "dict[str, Any]", - infer.predict(manifest_path, stages=stages, output_directory=out_dir), - ) + return infer.predict(manifest_path, stages=stages, output_directory=out_dir) def run_tutorial() -> dict[str, Any]: diff --git a/tutorials/tutorial_11_heart_and_lung_motion.py b/tutorials/tutorial_11_heart_and_lung_motion.py new file mode 100644 index 0000000..e6d5a6a --- /dev/null +++ b/tutorials/tutorial_11_heart_and_lung_motion.py @@ -0,0 +1,435 @@ +""" +Tutorial 11: Combined Heart and Lung Motion + +Purpose +------- +Build a single 4D animation of the DirLab ``Case1Pack`` thorax that combines +cardiac motion (a trained PhysicsNeMo MeshGraphNet applied to the patient heart) +with respiratory motion (the Tutorial 1 lung registration transforms). + +The script runs in two stages: + +1. **Cardiac motion.** Assemble the Tutorial 9 MeshGraphNet run directory, apply + the model to the Case1Pack heart fit (Tutorial 5) at each cardiac stage, and + rasterize each stage's displacement onto the Case1Pack grid as a deformation + field (``WorkflowInferPhysicsNeMoMGN.create_deformation_field``). A + beating-heart 4D USD of the heart surfaces is written for reference. + +2. **Combined motion.** Build a per-cell-labeled thorax surface by contouring the + Tutorial 2 patient labelmap, smooth each cardiac deformation field and apply it + at the reference frame, then carry the cardiac-deformed surface through the + respiratory cycle with the Tutorial 1 forward transforms. Breathing and the + heartbeat run as independent rhythms (see Composition order). The 100 combined + frames are written as VTP files and assembled into a single animated 4D USD + that is split by anatomy label and painted with per-organ OmniSurface + materials. + +Anatomy materials +----------------- +The per-cell ``boundary_labels`` produced by contouring the labelmap propagate +unchanged through every warp (each frame is a deep copy with only its points +moved; label-preserving decimation is applied if enabled). ``ConvertVTKToUSD``, +given ``segmenter.taxonomy.all_labels()`` and the segmenter, splits each frame +into per-organ prims, and ``USDAnatomyTools.enhance_meshes`` then binds the +matching OmniSurface material (diffuse color, subsurface scattering, etc.). + +Composition order +----------------- +Cardiac deformation is applied first, at the reference frame where the cardiac +fields are defined, giving one warped surface per (breath phase, cardiac stage). +Each rendered frame then bilinearly interpolates that precomputed grid: the +respiratory axis advances with the breath phase, while the cardiac axis advances +**independently and continuously** at ``cardiac_cycles_per_phase`` beats per +phase. A value below 1.0 therefore lets a single heartbeat carry across a phase +boundary (0.75 = each breath phase covers three-quarters of a beat). Both axes +wrap, so the sequence loops. + +Inputs (hard-coded near the top; edit for your layout) +------ +- MGN model run directory (Tutorial 9), epoch-300 checkpoint. +- PCA template volume (``pca-vol-kcl/pca_mean.vtu``). +- Case1Pack reference image (``data/DirLab-4DCT/Case1Pack_T70.mha``). +- Case1Pack heart fit (Tutorial 5): registered coefficients + volume mesh. +- Respiratory forward transforms (Tutorial 1 lung output). +- Segmented patient labelmap (Tutorial 2 lung output). + +Outputs (under ``output/tutorial_11_heart_and_lung``) +------- +- ``deformation_field_s.mha`` / ``surface_normal_field_s.mha`` - + per-stage cardiac fields on the Case1Pack grid. +- ``deformed_heart_surface_s.vtp`` + ``beating_heart.usd`` - the beating + heart alone. +- ``combined_frame_.vtp`` (``000..099``) + ``heart_and_lung_motion.usd`` - + the combined respiratory + cardiac 4D motion, painted with anatomy materials. + +Prerequisites +------------- +Run Tutorials 1 (lung), 2 (lung), 5 (Case1Pack with the pca-vol-kcl model) and 9 +(MGN) first. Requires the ``[physicsnemo]`` extra. +""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path +from typing import cast + +import itk +import numpy as np +import pyvista as pv + +from physiotwin4d import ( + ContourTools, + ConvertVTKToUSD, + ImageTools, + SegmentChestTotalSegmentator, + TransformTools, + USDAnatomyTools, + 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) + logger = logging.getLogger("tutorial_11_heart_and_lung_motion") + + repo_root = Path(__file__).resolve().parent.parent + 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 + # PCA template volume the model was trained on. + mgn_model_dir = tutorials_dir / "output" / "tutorial_09_byod_mgn_3" + mgn_epoch = 300 + pca_mean_volume = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + + # Case1Pack reference image (defines the cardiac deformation-field grid). + reference_image_file = repo_root / "data" / "DirLab-4DCT" / "Case1Pack_T70.mha" + + # 15-mode pca-vol-kcl fit of the Case1Pack heart (Tutorial 5, re-run). + tutorial_05_dir = tutorials_dir / "output" / "tutorial_05_heart_to_lung" + coefficients_file = ( + tutorial_05_dir / "tutorial_05_heart_to_lung_registered_coefficients.json" + ) + registered_mesh_file = ( + tutorial_05_dir / "tutorial_05_heart_to_lung_template_mesh_registered.vtu" + ) + + # Respiratory forward transforms (Tutorial 1 lung) and the segmented patient + # labelmap (Tutorial 2 lung). The labelmap yields a per-cell-labeled surface + # so the final USD can be split by anatomy and painted with organ materials. + respiratory_dir = tutorials_dir / "output" / "tutorial_01_lung" + patient_labelmap_file = ( + tutorials_dir / "output" / "tutorial_02_lung" / "patient_labelmap.mha" + ) + + output_dir = tutorials_dir / "output" / "tutorial_11_heart_and_lung" + + # ---- Parameters -------------------------------------------------------- + # Cardiac stages sampled over one heartbeat (fraction of the RR interval). + cardiac_stages = [round(0.1 * k, 2) for k in range(10)] + # Fraction of a cardiac cycle advanced per respiratory phase. The heart is + # decoupled from breathing, so a value < 1 means a single beat continues into + # the next phase (1.0 locks exactly one full beat to each phase). + cardiac_cycles_per_phase = 0.75 + # Gaussian sigma (mm) used to smooth the sparse cardiac deformation fields. + smoothing_sigma_mm = 10.0 + # One-time conditioning of the patient surface, reused for every frame. + surface_decimation_reduction = 0.0 # fraction of triangles to remove + surface_smoothing_iterations = 0 # Taubin (non-shrinking) iterations + # USD playback rate; 10 stages/second ~= one heartbeat per second. + frames_per_second = 10.0 + + output_dir.mkdir(parents=True, exist_ok=True) + + for required in ( + reference_image_file, + coefficients_file, + registered_mesh_file, + pca_mean_volume, + patient_labelmap_file, + ): + if not required.exists(): + raise FileNotFoundError(f"Required input not found: {required}") + + forward_transform_files = sorted(respiratory_dir.glob("slice_*_all_forward.hdf")) + if not forward_transform_files: + raise FileNotFoundError( + f"No respiratory forward transforms found in {respiratory_dir}. " + "Run tutorial_01_lung_gated_ct_to_usd.py first." + ) + + transform_tools = TransformTools() + + # ======================================================================== + # 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) + + cardiac_fields: list[itk.Image] = [] + heart_surfaces: list[pv.PolyData] = [] + for stage in cardiac_stages: + result = infer.create_deformation_field( + shape_parameters=coefficients_file, + stage=float(stage), + reference_image=reference_image, + reference_surface=registered_mesh_file, + ) + pct = int(round(stage * 100)) + itk.imwrite( + result["deformation_field"], + str(output_dir / f"deformation_field_s{pct:03d}.mha"), + compression=True, + ) + itk.imwrite( + result["normal_image"], + str(output_dir / f"surface_normal_field_s{pct:03d}.mha"), + compression=True, + ) + result["deformed_surface"].save( + str(output_dir / f"deformed_heart_surface_s{pct:03d}.vtp") + ) + cardiac_fields.append(result["deformation_field"]) + heart_surfaces.append(result["deformed_surface"]) + logger.info("cardiac stage %.2f -> deformation_field_s%03d.mha", stage, pct) + + # Beating-heart-only 4D USD (heart surfaces, one heartbeat per second). + heart_usd = WorkflowConvertVTKToUSD( + input_meshes=heart_surfaces, + usd_project_name="beating_heart", + output_directory=output_dir, + appearance="anatomy", + anatomy_type="heart", + separate_by_connectivity=True, + frames_per_second=float(len(cardiac_stages)), + ) + heart_usd.process() + + # ======================================================================== + # Stage 2: combine the cardiac fields with respiratory motion. + # ======================================================================== + cardiac_transforms = [ + _smoothed_cardiac_transform(field, smoothing_sigma_mm, transform_tools) + for field in cardiac_fields + ] + logger.info( + "Smoothed %d cardiac deformation fields by %.1f mm", + len(cardiac_transforms), + smoothing_sigma_mm, + ) + + # Labeled contour surface: each cell carries `boundary_labels`, which survive + # the warps (deep copies) and let ConvertVTKToUSD split the USD by anatomy so + # USDAnatomyTools.enhance_meshes can bind per-organ OmniSurface materials. + 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, surface_decimation_reduction, surface_smoothing_iterations + ) + logger.info( + "Patient surface: %d points, %d cells (decimation=%.2f, smoothing_iters=%d)", + patient_surface.n_points, + patient_surface.n_cells, + surface_decimation_reduction, + surface_smoothing_iterations, + ) + + # Cardiac motion at the reference frame, once per stage, reused across phases. + cardiac_surfaces = [ + transform_tools.transform_pvcontour(patient_surface, cardiac_transform) + for cardiac_transform in cardiac_transforms + ] + + n_phases = len(forward_transform_files) + n_stages = len(cardiac_surfaces) + + # 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)) + resp_points.append( + [ + np.asarray( + transform_tools.transform_pvcontour( + cardiac_surface, forward_transform + ).points, + dtype=np.float32, + ) + for cardiac_surface in cardiac_surfaces + ] + ) + 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() + + # Render frames by bilinearly interpolating the precomputed (phase, stage) + # warped-point grid: respiratory advances with the breath phase, while the + # cardiac cycle advances continuously and independently at + # ``cardiac_cycles_per_phase`` beats per phase, so a heartbeat carries across + # phase boundaries. Both axes wrap, so the sequence loops. + frames_per_phase = n_stages + n_frames = n_phases * frames_per_phase + combined_files: list[Path] = [] + usd_frames: list[pv.PolyData] = [] + for frame_idx in range(n_frames): + # Respiratory position: current breath phase and fraction into it. + phase_pos = frame_idx / frames_per_phase + phase_idx = int(phase_pos) + next_phase_idx = (phase_idx + 1) % n_phases + resp_blend = phase_pos - phase_idx + + # Cardiac position: continuous across phases, wrapping within the cycle. + card_pos = (phase_pos * cardiac_cycles_per_phase * n_stages) % n_stages + stage_idx = int(card_pos) + next_stage_idx = (stage_idx + 1) % n_stages + card_blend = card_pos - stage_idx + + # Interpolate the cardiac cycle within each bounding phase, then between + # the two phases. + phase_a = (1.0 - card_blend) * resp_points[phase_idx][stage_idx] + ( + card_blend * resp_points[phase_idx][next_stage_idx] + ) + phase_b = (1.0 - card_blend) * resp_points[next_phase_idx][stage_idx] + ( + card_blend * resp_points[next_phase_idx][next_stage_idx] + ) + points = (1.0 - resp_blend) * phase_a + resp_blend * phase_b + + combined_surface = patient_surface.copy(deep=True) + combined_surface.points = points + + frame_file = output_dir / f"combined_frame_{frame_idx:03d}.vtp" + combined_surface.save(str(frame_file)) + combined_files.append(frame_file) + usd_frames.append(combined_surface) + + del resp_points + logger.info( + "Wrote %d combined-motion surfaces to %s", len(combined_files), output_dir + ) + + # Assemble the ordered frames into a single animated 4D USD, split by anatomy + # label (via the mesh `boundary_labels`) so each organ becomes its own prim + # under /World/heart_and_lung_motion/{group}/{label}. enhance_meshes then + # binds the per-organ OmniSurface materials (colors, subsurface scatter, ...). + segmenter = SegmentChestTotalSegmentator() + converter = ConvertVTKToUSD( + "heart_and_lung_motion", + usd_frames, + segmenter.taxonomy.all_labels(), + segmenter=segmenter, + frames_per_second=frames_per_second, + log_level=logging.INFO, + ) + usd_file = output_dir / "heart_and_lung_motion.usd" + stage = converter.convert(str(usd_file)) + USDAnatomyTools(stage).enhance_meshes(segmenter) + stage.Save() + logger.info("Wrote 4D USD with anatomy materials: %s", usd_file) + + tutorial_results = { + "cardiac_field_count": len(cardiac_fields), + "combined_surfaces": combined_files, + "beating_heart_usd": str(output_dir / "beating_heart.usd"), + "usd_file": str(usd_file), + } From fba375553ef079e1b9c9eab355b485aae7950be5 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Mon, 20 Jul 2026 14:56:57 -0400 Subject: [PATCH 3/3] ENH: Code rabbit --- .../segment_chest_total_segmentator.py | 2 +- ...art_to_lung_fit_statistical_model_to_patient.py | 14 +++++++++++--- tutorials/tutorial_11_heart_and_lung_motion.py | 2 +- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/physiotwin4d/segment_chest_total_segmentator.py b/src/physiotwin4d/segment_chest_total_segmentator.py index e317f7f..1cb032b 100644 --- a/src/physiotwin4d/segment_chest_total_segmentator.py +++ b/src/physiotwin4d/segment_chest_total_segmentator.py @@ -66,7 +66,7 @@ def __init__(self, log_level: int | str = logging.INFO): """ super().__init__(log_level=log_level) - self.target_spacing = 0.0 + self.target_spacing = 1.0 # TotalSegmentator class indices, grouped by anatomy. for group_name, organs in ( diff --git a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py index 09eb083..9541910 100644 --- a/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py @@ -99,17 +99,25 @@ patient_image = itk.imread(str(patient_image_file)) if not (output_dir / f"{project_name}_patient_image.nii.gz").exists(): - itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz") + itk.imwrite( + patient_image, + output_dir / f"{project_name}_patient_image.nii.gz", + compression=True, + ) segmentation_result = segmentation_method.segment(patient_image) patient_labelmap = segmentation_result["labelmap"] itk.imwrite( - patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz" + patient_labelmap, + output_dir / f"{project_name}_patient_labelmap.nii.gz", + compression=True, ) heart_labelmap = segmentation_result["heart"] itk.imwrite( - heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz" + heart_labelmap, + output_dir / f"{project_name}_heart_labelmap.nii.gz", + compression=True, ) contour_tools = ContourTools() diff --git a/tutorials/tutorial_11_heart_and_lung_motion.py b/tutorials/tutorial_11_heart_and_lung_motion.py index e6d5a6a..385b427 100644 --- a/tutorials/tutorial_11_heart_and_lung_motion.py +++ b/tutorials/tutorial_11_heart_and_lung_motion.py @@ -196,7 +196,7 @@ def _condition_surface( # Cardiac MGN model (Tutorial 9): epoch-300 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 = 300 + mgn_epoch = 1500 pca_mean_volume = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") # Case1Pack reference image (defines the cardiac deformation-field grid).