Register respiratory CT phases and reconstruct a higher-resolution 4D volume series.
- - 07 -Build a surface PCA model from five lung lobes and fit it to all cases.
- - - + 08 -Register respiratory phases with ANTs+ICON and propagate fitted meshes.
- +Fit a PCA heart model to the reference phase and propagate it to every gated phase with ICON registration.
+ - + 09 -Train a PhysicsNeMo MLP to predict lung-lobe meshes at requested stages.
+Train a PhysicsNeMo MeshGraphNet (9a) or MLP (9b) to predict cardiac meshes at requested stages.
+ + 10 +Load a Tutorial 9 checkpoint and predict cardiac surfaces at gated phases or caller-specified stages.
+ + Recommended Run Order @@ -89,9 +89,12 @@ arguments. 3. Run Tutorial 3 after downloading KCL-Heart-Model. 4. Run Tutorial 4 after Tutorial 3 because it can consume the PCA model output. 5. Run Tutorial 6 after downloading DirLab-4DCT. -6. Run Tutorial 7 after downloading DirLab-4DCT. -7. Run Tutorial 8 after Tutorial 7 because it consumes fitted PCA meshes. -8. Run Tutorial 9 after Tutorial 8 because it trains from propagated meshes. +6. Run Tutorial 8 after preparing your own cardiac gated CT, labelmaps, KCL + volume PCA model, and ICON weights (bring-your-own-data; see the note below). +7. Run Tutorial 9a and/or 9b after Tutorial 8 because they train from its + fitted meshes. +8. Run Tutorial 10a and/or 10b after Tutorial 9a / 9b because they evaluate + the trained checkpoints. Tutorial 1: Heart-Gated CT to Animated USD ========================================== @@ -215,74 +218,93 @@ Outputs Registered respiratory phases, reconstructed high-resolution CT volumes, and preview screenshots. -Tutorial 7: DirLab Lung-Lobe PCA Model -====================================== +.. note:: + + Tutorials 8-10 form the cardiac mesh stage-prediction pipeline and are + **bring-your-own-data**: unlike Tutorials 1-6 they do not use the repository + ``data/`` directory or a downloadable sample. Their path constants point at a + local ``D:/PhysioMotion4D/`` cardiac layout (gated CT, labelmaps, the KCL + volume PCA model, and ICON weights); edit those constants to match your own + data. The former DirLab lung-lobe PCA tutorial (number 7) has been removed; + numbering continues at 8. + +Tutorial 8: Fit the Cardiac SSM and Propagate Through Gated Phases +================================================================== Script - ``tutorials/tutorial_07_dirlab_pca_model.py`` + ``tutorials/tutorial_08_cardiac_fit_model.py`` Workflow - ``WorkflowConvertImageToVTK``, ``WorkflowCreateStatisticalModel``, and - ``WorkflowFitStatisticalModelToPatient`` + ``WorkflowFitStatisticalModelToPatient`` (PCA registration) and + ``WorkflowReconstructHighres4DCT`` (ICON time-series registration) Dataset - DirLab-4DCT, downloaded manually. + Bring your own cardiac gated CT, labelmaps, KCL volume PCA model, and ICON + weights under ``D:/PhysioMotion4D/``. Run .. code-block:: bash - python tutorials/tutorial_07_dirlab_pca_model.py + python tutorials/tutorial_08_cardiac_fit_model.py Outputs - Five-lobe lung surface meshes, a surface PCA model, and PCA-fitted surfaces - for every available case. + Per-patient fitted SSM mesh/surface, PCA coefficients, and the SSM warped to + every gated phase, all written under ``OUTPUT_DIR``. -Tutorial 8: DirLab PCA Time-Series Propagation -============================================== +Tutorial 9a / 9b: Train a PhysicsNeMo Cardiac Stage Model +========================================================= Script - ``tutorials/tutorial_08_dirlab_pca_time_series.py`` + ``tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py`` (MLP) Workflow - ``RegisterTimeSeriesImages`` with ``registration_method='Greedy_ICON'`` and - ``TransformTools`` + ``physicsnemo.models.meshgraphnet.MeshGraphNet`` (9a) and + ``physicsnemo.models.mlp.FullyConnected`` (9b), trained on Tutorial 8 meshes. Dataset - DirLab-4DCT plus Tutorial 7 fitted mesh outputs. + Tutorial 8 fitted-mesh outputs. + +Extra install + PhysicsNeMo is an optional dependency. Install with + ``pip install "physiomotion4d[physicsnemo]"`` (requires Python >= 3.11). The + MeshGraphNet variant also requires ``torch-geometric``. Run .. code-block:: bash - python tutorials/tutorial_08_dirlab_pca_time_series.py + python tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py + python tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py Outputs - Per-case ANTs+ICON transforms and one PCA-fitted lung-lobe surface for each - DirLab respiratory phase. + Shared PhysicsNeMo checkpoints, training metadata, loss / RMSE histories, and + held-out predictions written under each trainer's ``OUTPUT_DIR``. -Tutorial 9: PhysicsNeMo Mesh Stage Model -======================================== +Tutorial 10a / 10b: Predict and Evaluate Cardiac Surfaces +========================================================= Script - ``tutorials/tutorial_09_physicsnemo_mesh_stage_model.py`` + ``tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py`` (MLP) Workflow - ``physicsnemo.models.mlp.FullyConnected`` trained on Tutorial 8 meshes. + Load a Tutorial 9 checkpoint and predict cardiac surfaces for one subject at + each gated phase (with error statistics) or at caller-specified stages. Dataset - Tutorial 8 propagated PCA mesh outputs. - -Extra install - PhysicsNeMo is an optional dependency. Install with - ``pip install "physiomotion4d[physicsnemo]"`` (requires Python >= 3.11). + Tutorial 9a / 9b trained checkpoints plus the Tutorial 8 fitted meshes. Run .. code-block:: bash - python tutorials/tutorial_09_physicsnemo_mesh_stage_model.py + python tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py pm0002 --epoch 5000 --out results/pm0002 + + Run with no arguments to use the ``run_tutorial`` entry point and its + ``DEFAULT_SUBJECT`` / ``DEFAULT_EPOCH`` constants. Outputs - Per-case PhysicsNeMo checkpoints, training metadata, loss histories, and a - predicted PCA-fitted mesh at the requested normalized respiratory stage. + Predicted ``.vtp`` surfaces per phase (with per-point error arrays when + ground truth exists) and a ``statistics.csv`` error summary. Dataset Notes ============= diff --git a/experiments/LongitudinalRegistration/.gitignore b/experiments/LongitudinalRegistration/.gitignore index f850328..969a911 100644 --- a/experiments/LongitudinalRegistration/.gitignore +++ b/experiments/LongitudinalRegistration/.gitignore @@ -1,2 +1,3 @@ uniGradICON fixed +*.txt diff --git a/pyproject.toml b/pyproject.toml index e43af6f..890e9d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,7 @@ dev = [ "pre-commit>=3.0.0", "pytest>=7.0.0", "pytest-cov>=4.0.0", - "mypy>=1.0.0", + "mypy>=2.1.0", "ruff>=0.1.0", "scipy-stubs>=1.14.0.0,<1.16.0.0", "types-requests>=2.32.0.0", @@ -268,7 +268,11 @@ module = [ "physiomotion4d.cli.visualize_pca_modes", "physiomotion4d.vtk_to_usd.mesh_utils", "physiomotion4d.vtk_to_usd.vtk_reader", - "tutorial_09_physicsnemo_mesh_stage_model", + "tutorial_08_cardiac_fit_model", + "tutorial_09a_cardiac_train_physicsnemo_mgn", + "tutorial_09b_cardiac_train_physicsnemo_mlp", + "tutorial_10a_cardiac_eval_physicsnemo_mgn", + "tutorial_10b_cardiac_eval_physicsnemo_mlp", ] disable_error_code = ["import-not-found", "import-untyped"] diff --git a/src/physiomotion4d/workflow_reconstruct_highres_4d_ct.py b/src/physiomotion4d/workflow_reconstruct_highres_4d_ct.py index 9c63d03..7cd6cc7 100644 --- a/src/physiomotion4d/workflow_reconstruct_highres_4d_ct.py +++ b/src/physiomotion4d/workflow_reconstruct_highres_4d_ct.py @@ -160,7 +160,7 @@ def __init__( # Set default number of iterations based on registration method self.number_of_iterations_Greedy: list[int] = [30, 15, 7, 3] - self.number_of_iterations_ICON: int = 20 + self.number_of_iterations_ICON: Optional[int] = 20 # Initialize registrar self.registrar = RegisterTimeSeriesImages( @@ -184,11 +184,15 @@ def set_number_of_iterations_Greedy( """ self.number_of_iterations_Greedy = number_of_iterations_Greedy - def set_number_of_iterations_ICON(self, number_of_iterations_ICON: int) -> None: + def set_number_of_iterations_ICON( + self, number_of_iterations_ICON: Optional[int] + ) -> None: """Set the number of iterations for ICON registration. Args: - number_of_iterations_ICON: Number of fine-tuning steps for ICON + number_of_iterations_ICON: Number of fine-tuning steps for ICON. + If None, ICON fine-tuning is disabled (the pretrained network + is used as-is). """ self.number_of_iterations_ICON = number_of_iterations_ICON diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py index cb4839b..9fcc372 100644 --- a/tests/test_tutorials.py +++ b/tests/test_tutorials.py @@ -21,7 +21,9 @@ from __future__ import annotations +import importlib.util import runpy +import sys from pathlib import Path from typing import Any @@ -231,6 +233,131 @@ def test_run(self, test_directories: dict[str, Path]) -> None: _compare_screenshots(results["screenshots"], tt) +# ----------------------------------------------------------------------------- +# Tutorials 8-10 - Cardiac mesh stage-prediction pipeline (bring-your-own-data) +# +# These tutorials use a local ``D:/PhysioMotion4D/`` cardiac dataset and (for +# Tutorials 9 and 10) the optional PhysicsNeMo dependency, so they are skipped +# automatically unless that data / those checkpoints are present. They produce +# no screenshots; the tests assert the tutorial ran and populated +# ``tutorial_results``. +# ----------------------------------------------------------------------------- + +_CARDIAC_DATA_ROOT = Path("D:/PhysioMotion4D") +_CARDIAC_FITTED_MESHES_DIR = _CARDIAC_DATA_ROOT / "duke_data" / "fitted_kcl_meshes" +_TUTORIALS_DIR = _REPO_ROOT / "tutorials" + + +def _physicsnemo_available() -> bool: + """True if the optional PhysicsNeMo dependency is importable.""" + return importlib.util.find_spec("physicsnemo") is not None + + +def _torch_geometric_available() -> bool: + """True if the optional PyTorch Geometric dependency is importable.""" + return importlib.util.find_spec("torch_geometric") is not None + + +def _run_eval_tutorial(script_name: str) -> dict[str, Any]: + """Run a Tutorial 10 eval script through its no-argument ``run_tutorial`` path. + + ``runpy`` does not reset ``sys.argv``, so without this the eval scripts would + see pytest's arguments and try to parse them as CLI options. Force a single + argv entry so the ``len(sys.argv) > 1`` dispatch selects ``run_tutorial``. + """ + saved_argv = sys.argv + sys.argv = [script_name] + try: + return _run_tutorial_script(script_name) + finally: + sys.argv = saved_argv + + +@pytest.mark.tutorial +@pytest.mark.slow +class TestTutorial08CardiacFitModel: + """End-to-end test for tutorial_08_cardiac_fit_model.py (bring-your-own-data).""" + + def test_run(self) -> None: + if not (_CARDIAC_DATA_ROOT / "duke_data" / "gated_nii").exists(): + pytest.skip( + "Cardiac dataset not present at D:/PhysioMotion4D/. Tutorial 8 is " + "bring-your-own-data; see tutorials/README.md." + ) + results = _run_tutorial_script("tutorial_08_cardiac_fit_model.py") + assert "patients" in results, "Tutorial 8 should report processed patients" + + +@pytest.mark.tutorial +@pytest.mark.slow +@pytest.mark.requires_gpu +class TestTutorial09aCardiacTrainMGN: + """End-to-end test for tutorial_09a_cardiac_train_physicsnemo_mgn.py.""" + + def test_run(self) -> None: + if not _physicsnemo_available(): + pytest.skip("PhysicsNeMo not installed (optional [physicsnemo] extra).") + if not _torch_geometric_available(): + pytest.skip("torch-geometric not installed (required for MeshGraphNet).") + if not _CARDIAC_FITTED_MESHES_DIR.exists(): + pytest.skip("Tutorial 8 cardiac output not present; run Tutorial 8 first.") + results = _run_tutorial_script("tutorial_09a_cardiac_train_physicsnemo_mgn.py") + assert isinstance(results, dict) + + +@pytest.mark.tutorial +@pytest.mark.slow +@pytest.mark.requires_gpu +class TestTutorial09bCardiacTrainMLP: + """End-to-end test for tutorial_09b_cardiac_train_physicsnemo_mlp.py.""" + + def test_run(self) -> None: + if not _physicsnemo_available(): + pytest.skip("PhysicsNeMo not installed (optional [physicsnemo] extra).") + if not _CARDIAC_FITTED_MESHES_DIR.exists(): + pytest.skip("Tutorial 8 cardiac output not present; run Tutorial 8 first.") + results = _run_tutorial_script("tutorial_09b_cardiac_train_physicsnemo_mlp.py") + assert isinstance(results, dict) + + +@pytest.mark.tutorial +@pytest.mark.slow +class TestTutorial10aCardiacEvalMGN: + """End-to-end test for tutorial_10a_cardiac_eval_physicsnemo_mgn.py.""" + + def test_run(self) -> None: + if not _physicsnemo_available(): + pytest.skip("PhysicsNeMo not installed (optional [physicsnemo] extra).") + if not _torch_geometric_available(): + pytest.skip("torch-geometric not installed (required for MeshGraphNet).") + checkpoint = _TUTORIALS_DIR / "output_mgn" / "mgn_stage_model.pt" + if not checkpoint.exists() or not _CARDIAC_FITTED_MESHES_DIR.exists(): + pytest.skip( + "Tutorial 9a checkpoint or cardiac data not present; " + "run Tutorials 8 and 9a first." + ) + results = _run_eval_tutorial("tutorial_10a_cardiac_eval_physicsnemo_mgn.py") + assert "predicted_files" in results + + +@pytest.mark.tutorial +@pytest.mark.slow +class TestTutorial10bCardiacEvalMLP: + """End-to-end test for tutorial_10b_cardiac_eval_physicsnemo_mlp.py.""" + + def test_run(self) -> None: + if not _physicsnemo_available(): + pytest.skip("PhysicsNeMo not installed (optional [physicsnemo] extra).") + checkpoint = _TUTORIALS_DIR / "output" / "physicsnemo_stage_model.pt" + if not checkpoint.exists() or not _CARDIAC_FITTED_MESHES_DIR.exists(): + pytest.skip( + "Tutorial 9b checkpoint or cardiac data not present; " + "run Tutorials 8 and 9b first." + ) + results = _run_eval_tutorial("tutorial_10b_cardiac_eval_physicsnemo_mlp.py") + assert "predicted_files" in results + + # ----------------------------------------------------------------------------- # Tutorial 6 - Reconstruct High-Resolution 4D CT # ----------------------------------------------------------------------------- diff --git a/tutorials/README.md b/tutorials/README.md index 6357e2b..656dd2c 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -19,9 +19,18 @@ dataset licensing, and expected directory layout. | 4 | [tutorial_04_fit_statistical_model_to_patient.py](tutorial_04_fit_statistical_model_to_patient.py) | `WorkflowFitStatisticalModelToPatient` | KCL-Heart-Model plus Tutorial 3 output | | 5 | [tutorial_05_vtk_to_usd.py](tutorial_05_vtk_to_usd.py) | `WorkflowConvertVTKToUSD` | Output of tutorial 2 | | 6 | [tutorial_06_reconstruct_highres_4d_ct.py](tutorial_06_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | -| 7 | [tutorial_07_dirlab_pca_model.py](tutorial_07_dirlab_pca_model.py) | `WorkflowCreateStatisticalModel`, `WorkflowFitStatisticalModelToPatient` | DirLab-4DCT (manual) | -| 8 | [tutorial_08_dirlab_pca_time_series.py](tutorial_08_dirlab_pca_time_series.py) | `RegisterTimeSeriesImages` | DirLab-4DCT plus Tutorial 7 output | -| 9 | [tutorial_09_physicsnemo_mesh_stage_model.py](tutorial_09_physicsnemo_mesh_stage_model.py) | `physicsnemo.models.mlp.FullyConnected` (requires `[physicsnemo]` extra) | Tutorial 8 output | +| 8 | [tutorial_08_cardiac_fit_model.py](tutorial_08_cardiac_fit_model.py) | `WorkflowFitStatisticalModelToPatient`, `WorkflowReconstructHighres4DCT` | Bring your own (cardiac gated CT, `D:/PhysioMotion4D/`) | +| 9a | [tutorial_09a_cardiac_train_physicsnemo_mgn.py](tutorial_09a_cardiac_train_physicsnemo_mgn.py) | `physicsnemo.models.meshgraphnet.MeshGraphNet` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 8 output | +| 9b | [tutorial_09b_cardiac_train_physicsnemo_mlp.py](tutorial_09b_cardiac_train_physicsnemo_mlp.py) | `physicsnemo.models.mlp.FullyConnected` (requires `[physicsnemo]` extra) | Tutorial 8 output | +| 10a | [tutorial_10a_cardiac_eval_physicsnemo_mgn.py](tutorial_10a_cardiac_eval_physicsnemo_mgn.py) | `physicsnemo.models.meshgraphnet.MeshGraphNet` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 9a checkpoint | +| 10b | [tutorial_10b_cardiac_eval_physicsnemo_mlp.py](tutorial_10b_cardiac_eval_physicsnemo_mlp.py) | `physicsnemo.models.mlp.FullyConnected` (requires `[physicsnemo]` extra) | Tutorial 9b checkpoint | + +> **Tutorials 8-10 are bring-your-own-data.** Unlike Tutorials 1-6, they do not +> use the repository `data/` directory or a downloadable sample. Their path +> constants point at a local `D:/PhysioMotion4D/` cardiac layout (gated CT, +> labelmaps, the KCL volume PCA model, and ICON weights); edit those constants to +> match your own data. (The former DirLab lung-lobe PCA tutorial, number 7, has +> been removed; the numbering continues at 8.) ## Running a Tutorial @@ -70,9 +79,13 @@ pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD --run-tutorials 3. **Tutorial 3** creates the PCA statistical model from KCL-Heart-Model. 4. **Tutorial 4** applies the statistical model, consuming Tutorial 3 output. 5. **Tutorial 6** requires DirLab-4DCT - download it per `data/README.md`. -6. **Tutorial 7** creates a surface PCA model of the five lung lobes from DirLab-4DCT, then fits it to every available case. -7. **Tutorial 8** registers DirLab respiratory phases with ANTs+ICON and propagates the Tutorial 7 fitted meshes through each time series. -8. **Tutorial 9** trains a PhysicsNeMo model to predict a PCA-fitted mesh at a user-specified respiratory stage. PhysicsNeMo is an optional extra: install with `pip install "physiomotion4d[physicsnemo]"` (requires Python >= 3.11). + +The cardiac mesh stage-prediction pipeline (Tutorials 8 -> 9 -> 10) is +bring-your-own-data and runs in order: + +6. **Tutorial 8** fits the KCL cardiac PCA model to each patient's reference CT and propagates the fitted SSM mesh through every gated phase (output feeds Tutorial 9). +7. **Tutorial 9a / 9b** train a PhysicsNeMo MeshGraphNet (9a) and MLP (9b) to predict a cardiac surface at any cardiac stage. PhysicsNeMo is an optional extra: install with `pip install "physiomotion4d[physicsnemo]"` (requires Python >= 3.11); the MeshGraphNet also needs `torch-geometric`. +8. **Tutorial 10a / 10b** load a trained MeshGraphNet (10a) or MLP (10b) checkpoint and predict / score cardiac surfaces for one subject. Each can be run from the command line or, with no arguments, via its `run_tutorial` entry point. ## For Contributors diff --git a/tutorials/tutorial_04_fit_statistical_model_to_patient.py b/tutorials/tutorial_04_fit_statistical_model_to_patient.py index 810ce11..ae96b47 100644 --- a/tutorials/tutorial_04_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_04_fit_statistical_model_to_patient.py @@ -20,7 +20,7 @@ import json import logging from pathlib import Path -from typing import Any, cast +from typing import Any, Optional, cast import pyvista as pv @@ -69,10 +69,7 @@ if isinstance(template_data, pv.PolyData): template_model = template_data else: - template_model = cast( - pv.PolyData, - template_data.extract_surface(algorithm="dataset_surface"), - ) + template_model = template_data.extract_surface(algorithm="dataset_surface") sample_files = sorted((data_dir / "sample_meshes").glob("*.vtu")) if not sample_files: @@ -85,20 +82,17 @@ "See data/README.md for download instructions." ) - patient_models: list[pv.PolyData] = [] + patient_models: list[pv.DataSet] = [] for sample_file in sample_files: sample_data = cast(pv.DataSet, pv.read(str(sample_file))) if isinstance(sample_data, pv.PolyData): patient_models.append(sample_data) else: patient_models.append( - cast( - pv.PolyData, - sample_data.extract_surface(algorithm="dataset_surface"), - ) + sample_data.extract_surface(algorithm="dataset_surface") ) - pca_model: dict[str, Any] | None = None + pca_model: Optional[dict[str, Any]] = None if pca_json.exists(): with pca_json.open(encoding="utf-8") as f: pca_model = json.load(f) diff --git a/tutorials/tutorial_07_dirlab_pca_model.py b/tutorials/tutorial_07_dirlab_pca_model.py deleted file mode 100644 index fdd746a..0000000 --- a/tutorials/tutorial_07_dirlab_pca_model.py +++ /dev/null @@ -1,305 +0,0 @@ -"""Tutorial 7: Build and fit a PCA lung-lobe model from DirLab 4D CT cases. - -This tutorial uses one respiratory phase from each available DirLab case. It -segments the five lung lobes, builds a surface PCA model from all but two -cases, then fits that model to every available case. -""" - -# %% -from __future__ import annotations - -import json -import logging -from pathlib import Path -from typing import Any, cast - -import itk -import numpy as np -import pyvista as pv - -from physiomotion4d.contour_tools import ContourTools -from physiomotion4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK -from physiomotion4d.workflow_create_statistical_model import ( - WorkflowCreateStatisticalModel, -) -from physiomotion4d.test_tools import TestTools -from physiomotion4d.workflow_fit_statistical_model_to_patient import ( - WorkflowFitStatisticalModelToPatient, -) - - -# 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. Wrapping consistently across -# tutorials also matches the style of tutorial_01. -if __name__ == "__main__": - # %% - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - DATA_DIR = REPO_ROOT / "data" - FULL_DATA_DIR = DATA_DIR - TEST_DATA_DIR = DATA_DIR / "test" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_07_dirlab_pca_model" - PCA_COMPONENTS = 5 - LOG_LEVEL = logging.INFO - - # %% - def create_meshes( - data_dir: Path, - output_dir: Path, - log_level: int = logging.INFO, - ) -> dict[str, Path]: - """Segment DirLab CT cases and save one five-lobe surface mesh per case. - - Parameters - ---------- - data_dir - Repository data directory containing ``DirLab-4DCT``. - output_dir - Directory where tutorial outputs are written. - log_level - Logging level for the conversion workflow. - - Returns - ------- - dict[str, Path] - Output directory and saved lung-lobe surface mesh filenames. - """ - - dirlab_dir = data_dir / "DirLab-4DCT" - meshes_dir = output_dir / "meshes" - meshes_dir.mkdir(parents=True, exist_ok=True) - - case_prefixes = [ - "Case1Pack", - "Case2Pack", - "Case3Pack", - "Case4Pack", - "Case5Pack", - "Case6Pack", - "Case7Pack", - "Case8Deploy", - "Case9Pack", - "Case10Pack", - ] - lung_lobe_ids = { - 10: "lung_upper_lobe_left", - 11: "lung_lower_lobe_left", - 12: "lung_upper_lobe_right", - 13: "lung_middle_lobe_right", - 14: "lung_lower_lobe_right", - } - - mesh_files: list[Path] = [] - contour_tools = ContourTools(log_level=log_level) - for case_number, case_prefix in enumerate(case_prefixes, start=1): - mesh_file = meshes_dir / f"{case_prefix}_lung_lobes.vtp" - if mesh_file.exists(): - mesh_files.append(mesh_file) - continue - - case_dir = dirlab_dir / f"Case{case_number}" - phase_files = sorted(case_dir.glob("*.mha")) + sorted( - case_dir.glob("*.mhd") - ) - if not phase_files: - phase_files = sorted(dirlab_dir.glob(f"{case_prefix}_T*.mha")) - phase_files += sorted(dirlab_dir.glob(f"{case_prefix}_T*.mhd")) - if not phase_files and case_number == 8: - phase_files = sorted(dirlab_dir.glob("Case8Pack_T*.mha")) - phase_files += sorted(dirlab_dir.glob("Case8Pack_T*.mhd")) - if not phase_files: - print(f"Skipping {case_prefix}: no DirLab phase image found") - continue - - print(f"Segmenting {case_prefix} from {phase_files[0].name}") - image = itk.imread(str(phase_files[0])) - workflow = WorkflowConvertImageToVTK( - segmentation_method="ChestTotalSegmentator", - log_level=log_level, - ) - result = workflow.run_workflow( - input_image=image, - contrast_enhanced_study=False, - anatomy_groups=["lung"], - ) - - labelmap = result["labelmap"] - labelmap_arr = itk.GetArrayFromImage(labelmap) - lobe_surfaces: list[pv.PolyData] = [] - for label_id, lobe_name in lung_lobe_ids.items(): - lobe_arr = (labelmap_arr == label_id).astype(np.uint8) - if int(lobe_arr.sum()) == 0: - print(f"Skipping {case_prefix}: missing {lobe_name}") - lobe_surfaces = [] - break - - lobe_mask = itk.GetImageFromArray(lobe_arr) - lobe_mask.CopyInformation(labelmap) - lobe_surface = contour_tools.extract_contours(lobe_mask) - lobe_surface.field_data["LungLobeName"] = np.array([lobe_name]) - lobe_surface.field_data["LungLobeLabel"] = np.array([label_id]) - lobe_surfaces.append(lobe_surface) - - if len(lobe_surfaces) != len(lung_lobe_ids): - continue - - lung_surface = cast( - pv.PolyData, - pv.merge(lobe_surfaces, merge_points=False), - ) - lung_surface.save(mesh_file) - mesh_files.append(mesh_file) - - return {"meshes_dir": meshes_dir, **{mesh.stem: mesh for mesh in mesh_files}} - - def create_model( - meshes_dir: Path, - output_dir: Path, - pca_components: int = 5, - log_level: int = logging.INFO, - ) -> dict[str, Any]: - """Create a surface PCA model from all but the final two DirLab lobe meshes. - - Parameters - ---------- - meshes_dir - Directory containing lung-lobe surface files from :func:`create_meshes`. - output_dir - Directory where model outputs are written. - pca_components - Requested PCA component count. - log_level - Logging level for the model creation workflow. - - Returns - ------- - dict[str, Any] - PCA model, saved model filenames, training files, and all mesh filenames. - """ - - model_dir = output_dir / "pca_model" - model_dir.mkdir(parents=True, exist_ok=True) - - mesh_files = sorted(meshes_dir.glob("*_lung_lobes.vtp")) - if len(mesh_files) < 4: - raise ValueError("At least four DirLab lung-lobe meshes are needed.") - - training_files = mesh_files[:-2] - held_out_files = mesh_files[-2:] - sample_meshes = [pv.read(str(mesh_file)) for mesh_file in training_files] - reference_mesh = pv.read(str(training_files[0])) - component_count = min(pca_components, max(1, len(sample_meshes) - 1)) - - workflow = WorkflowCreateStatisticalModel( - sample_meshes=sample_meshes, - reference_mesh=reference_mesh, - pca_number_of_components=component_count, - solve_for_surface_pca=True, - log_level=log_level, - ) - result = workflow.run_workflow() - - mean_surface_file = model_dir / "pca_mean_surface.vtp" - pca_model_file = model_dir / "pca_model.json" - result["pca_mean_surface"].save(mean_surface_file) - pca_model_file.write_text( - json.dumps(result["pca_model"], indent=2), - encoding="utf-8", - ) - - return { - "pca_model": result["pca_model"], - "mean_surface_file": mean_surface_file, - "pca_model_file": pca_model_file, - "training_files": training_files, - "held_out_files": held_out_files, - "mesh_files": mesh_files, - } - - def fit_model( - model_result: dict[str, Any], - output_dir: Path, - log_level: int = logging.INFO, - ) -> dict[str, Path]: - """Fit the surface PCA lung-lobe model to every available DirLab mesh. - - Parameters - ---------- - model_result - Dictionary returned by :func:`create_model`. - output_dir - Directory where fit outputs are written. - log_level - Logging level for the fitting workflow. - - Returns - ------- - dict[str, Path] - Saved fitted surface filenames. - """ - - fits_dir = output_dir / "fits" - fits_dir.mkdir(parents=True, exist_ok=True) - - template_model = pv.read(str(model_result["mean_surface_file"])) - fitted_files: dict[str, Path] = {} - - for patient_file in model_result["mesh_files"]: - patient_model = pv.read(str(patient_file)) - workflow = WorkflowFitStatisticalModelToPatient( - template_model=template_model, - patient_models=[patient_model], - log_level=log_level, - ) - workflow.set_use_pca_registration( - True, - pca_model=model_result["pca_model"], - pca_number_of_modes=0, - pca_uses_surface=True, - ) - workflow.set_use_labelmap_to_labelmap_registration(False) - - result = workflow.run_workflow() - fitted_surface = result["registered_template_model_surface"] - fitted_file = fits_dir / f"{patient_file.stem}_pca_fit.vtp" - fitted_surface.save(fitted_file) - fitted_files[patient_file.stem] = fitted_file - - return fitted_files - - def run_tutorial() -> dict[str, Any]: - """Run mesh creation, PCA model creation, and PCA fitting in sequence. - - Returns - ------- - dict[str, Any] - Mesh, PCA model, and fitted model output information. - """ - - data_dir = TEST_DATA_DIR if TestTools.running_as_test() else FULL_DATA_DIR - output_dir = OUTPUT_DIR - pca_components = PCA_COMPONENTS - log_level = LOG_LEVEL - - output_dir.mkdir(parents=True, exist_ok=True) - mesh_result = create_meshes(data_dir, output_dir, log_level=log_level) - model_result = create_model( - mesh_result["meshes_dir"], - output_dir, - pca_components=pca_components, - log_level=log_level, - ) - fit_result = fit_model(model_result, output_dir, log_level=log_level) - - return { - "mesh_result": mesh_result, - "model_result": model_result, - "fit_result": fit_result, - } - - # %% - # Run this cell in VS Code or Cursor: - tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_08_cardiac_fit_model.py b/tutorials/tutorial_08_cardiac_fit_model.py new file mode 100644 index 0000000..8737ae3 --- /dev/null +++ b/tutorials/tutorial_08_cardiac_fit_model.py @@ -0,0 +1,297 @@ +""" +Tutorial 8: Fit the Cardiac SSM and Propagate It Through Gated Phases + +Purpose +------- +First stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +For each patient it turns gated CT scans into the statistical-shape-model (SSM) +surfaces and volume meshes that the Tutorial 9 trainers +(``tutorial_09a_cardiac_train_physicsnemo_mgn.py`` / +``tutorial_09b_cardiac_train_physicsnemo_mlp.py``) consume: + +1. Fit the KCL PCA heart model to the reference phase. A surface is extracted + from the reference labelmap and the KCL PCA volume model is fitted with + PCA-based registration (``WorkflowFitStatisticalModelToPatient`` with + ``use_pca_registration=True``, surface fitting disabled). This yields the + patient's PCA coefficients plus the fitted SSM volume mesh and surface, all + sharing the model's fixed topology. + +2. Propagate the fitted mesh to every gated phase. Each gated time point is + registered to the reference with the deep-learning ICON registrar + (``WorkflowReconstructHighres4DCT``). The forward transform for each phase + warps the fitted SSM mesh and surface (``TransformTools.transform_pvcontour``, + with deformation magnitude attached), producing one + ``*_g{TT}_ssm_mesh.vtu`` / ``*_ssm_surface.vtp`` per phase. + +Bring Your Own Data +------------------- +This is a bring-your-own-data tutorial. Unlike Tutorials 1-6, it does not use the +repository ``data/`` directory or a downloadable sample; the path constants below +point at a local ``D:/PhysioMotion4D/`` layout. Edit them to match your own data. + +Data Required +------------- + * ``D:/PhysioMotion4D/duke_data/gated_nii/pm00??/`` - gated NIfTI CT per patient + * ``D:/PhysioMotion4D/duke_data/simple_ascardio/pm00??/`` - matching labelmaps + * ``D:/PhysioMotion4D/kcl-heart-pca/pca-vol-kcl/`` - PCA model (pca_mean.vtu, pca_model.json) + * ``D:/PhysioMotion4D/duke_data/icon_registration/`` - ICON registration weights + +Outputs (per patient, under ``OUTPUT_DIR/pm00??/``) +--------------------------------------------------- + * ``*_ssm_pca_coefficients.json`` - fitted PCA coefficient vector + * ``*_ssm_pca_mesh.vtu`` / ``*_ssm_pca_surface.vtp`` - PCA template before final warp + * ``*_ssm_mesh.vtu`` / ``*_ssm_surface.vtp`` - fitted reference SSM mesh/surface + * ``*_g{TT}_ssm_mesh.vtu`` / ``*_g{TT}_ssm_surface.vtp`` - SSM warped to gated phase TT% + * ``*_g{TT}_warped_ref.mha``, ``*_g{TT}_*_tfm.hdf`` - registration artifacts + * ``*_g{TT}_ref_labelmap.nii.gz`` - reference labelmap warped to each phase +""" + +# %% +# Imports +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any + +import itk +import pyvista as pv + +from physiomotion4d import ( + ContourTools, + TransformTools, + WorkflowFitStatisticalModelToPatient, + WorkflowReconstructHighres4DCT, +) +from physiomotion4d.test_tools import TestTools + +# 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__": + # %% + # Path configuration (bring-your-own-data: edit for your local layout) + DATA_DIR = Path("D:/PhysioMotion4D/duke_data/gated_nii") + LABELMAP_DIR = Path("D:/PhysioMotion4D/duke_data/simple_ascardio") + SSM_MEAN_MESH_FILE = Path( + "D:/PhysioMotion4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu" + ) + SSM_MODEL_FILE = Path("D:/PhysioMotion4D/kcl-heart-pca/pca-vol-kcl/pca_model.json") + ICON_WEIGHTS_PATH = Path( + "D:/PhysioMotion4D/duke_data/icon_registration/" + "icon_ct_cardiac_gated_weights.trch" + ) + # All outputs (fitted meshes, transforms, warped labelmaps) are written here; + # this is also the directory the Tutorial 9 trainers read from. + OUTPUT_DIR = Path("D:/PhysioMotion4D/duke_data/fitted_kcl_meshes") + # Simpleware's heart interior chamber labels, excluded from the distance map. + LABELMAP_INTERIOR_OBJECT_IDS = [1, 2, 3, 4] + # Recompute the expensive fit/registration steps (True) or reload cached + # results from OUTPUT_DIR (False). + RECOMPUTE = True + LOG_LEVEL = logging.INFO + + logging.basicConfig(level=LOG_LEVEL) + logger = logging.getLogger("tutorial_08_cardiac_fit_model") + + # In test mode, limit the run to a single patient to keep it tractable. + test_mode = TestTools.running_as_test() + + # %% + # Load the statistical atlas model + ssm_mean_mesh = pv.read(str(SSM_MEAN_MESH_FILE)) + with SSM_MODEL_FILE.open(encoding="utf-8") as f: + ssm_model = json.load(f) + + # %% + # Discover patients + patient_dirs = sorted(DATA_DIR.glob("pm00??")) + if test_mode: + patient_dirs = patient_dirs[:1] + + tutorial_results: dict[str, Any] = {"patients": {}} + + for patient_dir in patient_dirs: + patient_id = patient_dir.name + logger.info("%s", "=" * 48) + logger.info("Processing patient %s", patient_id) + logger.info("%s", "=" * 48) + + patient_output_dir = OUTPUT_DIR / patient_id + patient_output_dir.mkdir(parents=True, exist_ok=True) + + ref_image_files = list(patient_dir.glob("*ref.nii.gz")) + if len(ref_image_files) != 1: + raise ValueError(f"Expected 1 ref image file, found {len(ref_image_files)}") + ref_image_file = ref_image_files[0] + ref_image = itk.imread(str(ref_image_file)) + + ref_labelmap_file = ref_image_file.name.replace(".nii.gz", "_labelmap.nii.gz") + ref_labelmap = itk.imread(str(LABELMAP_DIR / patient_id / ref_labelmap_file)) + + # %% + # Step 1: fit the statistical model to the reference phase + contour_tools = ContourTools() + ref_surface = contour_tools.extract_contours(ref_labelmap) + + ssm_pca_coefficients_path = ( + patient_output_dir / f"{patient_id}_ssm_pca_coefficients.json" + ) + ssm_mesh_path = patient_output_dir / f"{patient_id}_ssm_mesh.vtu" + ssm_surface_path = patient_output_dir / f"{patient_id}_ssm_surface.vtp" + + if RECOMPUTE: + ssm_fit_workflow = WorkflowFitStatisticalModelToPatient( + template_model=ssm_mean_mesh, + patient_image=ref_image, + patient_models=[ref_surface], + patient_labelmap=ref_labelmap, + labelmap_interior_object_ids=LABELMAP_INTERIOR_OBJECT_IDS, + log_level=LOG_LEVEL, + ) + ssm_fit_workflow.set_use_pca_registration( + use_pca_registration=True, + pca_model=ssm_model, + pca_uses_surface=False, + ) + + ssm_fit_workflow_result = ssm_fit_workflow.run_workflow() + + ssm_pca_coefficients = ssm_fit_workflow.pca_coefficients + assert ssm_pca_coefficients is not None, ( + "pca_coefficients must be set after run_workflow() with " + "use_pca_registration=True" + ) + with ssm_pca_coefficients_path.open(mode="w", encoding="utf-8") as f: + json.dump(ssm_pca_coefficients.tolist(), f) + + ssm_pca_template_model = ssm_fit_workflow.pca_template_model + assert ssm_pca_template_model is not None + ssm_pca_template_model.save( + str(patient_output_dir / f"{patient_id}_ssm_pca_mesh.vtu") + ) + + ssm_pca_template_model_surface = ssm_fit_workflow.pca_template_model_surface + assert ssm_pca_template_model_surface is not None + ssm_pca_template_model_surface.save( + str(patient_output_dir / f"{patient_id}_ssm_pca_surface.vtp") + ) + + ssm_mesh_fitted = ssm_fit_workflow_result["registered_template_model"] + ssm_surface_fitted = ssm_fit_workflow_result[ + "registered_template_model_surface" + ] + + ssm_mesh_fitted.save(str(ssm_mesh_path)) + ssm_surface_fitted.save(str(ssm_surface_path)) + else: + ssm_mesh_fitted = pv.read(str(ssm_mesh_path)) + ssm_surface_fitted = pv.read(str(ssm_surface_path)) + + # %% + # Step 2: register every gated phase to the reference + gated_files = sorted( + file + for file in patient_dir.glob("*.nii.gz") + if file != ref_image_file and "nop" not in file.name and "_g" in file.stem + ) + + time_series = [] + time_series_ids = [] + for gated_file in gated_files: + time_series.append(itk.imread(str(gated_file))) + time_id = gated_file.name.split("_g")[1][:3] + time_series_ids.append(time_id) + + if RECOMPUTE: + reg_workflow = WorkflowReconstructHighres4DCT( + time_series_images=time_series, + fixed_image=ref_image, + registration_method="ICON", + ) + reg_workflow.registrar.registrar_ICON.set_weights_path( + str(ICON_WEIGHTS_PATH) + ) + reg_workflow.set_modality("ct") + reg_workflow.set_number_of_iterations_ICON(None) + reg_result = reg_workflow.run_workflow() + + reconstructed_images = reg_result["reconstructed_images"] + else: + reconstructed_images = [] + for time_id in time_series_ids: + image_path = ( + patient_output_dir / f"{patient_id}_g{time_id}_warped_ref.mha" + ) + reconstructed_images.append(itk.imread(str(image_path))) + + # %% + # Step 3: warp the fitted SSM mesh/surface to every gated phase + phase_outputs = [] + for image_index, image in enumerate(reconstructed_images): + time_id = time_series_ids[image_index] + logger.info("Patient %s: warping to time point %s", patient_id, time_id) + + if RECOMPUTE: + image_path = ( + patient_output_dir / f"{patient_id}_g{time_id}_warped_ref.mha" + ) + itk.imwrite(image, str(image_path), compression=True) + + fwd_tfm = reg_result["forward_transforms"][image_index] + itk.transformwrite( + fwd_tfm, + str( + patient_output_dir / f"{patient_id}_g{time_id}_forward_tfm.hdf" + ), + ) + + inv_tfm = reg_result["inverse_transforms"][image_index] + itk.transformwrite( + inv_tfm, + str( + patient_output_dir / f"{patient_id}_g{time_id}_inverse_tfm.hdf" + ), + ) + + # Warp the reference labelmap to this phase. Written under + # OUTPUT_DIR (never back into the input labelmap directory). + labelmap = TransformTools().transform_image( + ref_labelmap, inv_tfm, image, "nearest" + ) + itk.imwrite( + labelmap, + str( + patient_output_dir + / f"{patient_id}_g{time_id}_ref_labelmap.nii.gz" + ), + compression=True, + ) + else: + fwd_tfm = itk.transformread( + str(patient_output_dir / f"{patient_id}_g{time_id}_forward_tfm.hdf") + ) + + mesh = TransformTools().transform_pvcontour( + ssm_mesh_fitted, fwd_tfm, with_deformation_magnitude=True + ) + mesh_path = patient_output_dir / f"{patient_id}_g{time_id}_ssm_mesh.vtu" + mesh.save(str(mesh_path)) + + surface = TransformTools().transform_pvcontour( + ssm_surface_fitted, fwd_tfm, with_deformation_magnitude=True + ) + surface_path = ( + patient_output_dir / f"{patient_id}_g{time_id}_ssm_surface.vtp" + ) + surface.save(str(surface_path)) + phase_outputs.append({"time_id": time_id, "surface_file": surface_path}) + + tutorial_results["patients"][patient_id] = { + "pca_coefficients_file": ssm_pca_coefficients_path, + "ssm_surface_file": ssm_surface_path, + "phase_outputs": phase_outputs, + } diff --git a/tutorials/tutorial_08_dirlab_pca_time_series.py b/tutorials/tutorial_08_dirlab_pca_time_series.py deleted file mode 100644 index 55302ee..0000000 --- a/tutorials/tutorial_08_dirlab_pca_time_series.py +++ /dev/null @@ -1,206 +0,0 @@ -""" -Tutorial 8: Propagate DirLab PCA lung-lobe fits through each 4D CT time series. - -This tutorial uses the per-case PCA-fitted reference meshes created by Tutorial -7. For each DirLab case, it registers every respiratory phase to the reference -phase used by Tutorial 7, saves the image transforms, and applies the -reference-to-phase transform to the fitted mesh so each time point has a -PCA-derived lung-lobe surface. - -Data Required -------------- -See data/README.md for DirLab-4DCT download instructions. Run Tutorial 7 first -so ``output/tutorial_07_dirlab_pca_model/fits`` contains the per-case fitted -reference-stage meshes. -""" - -# %% -from __future__ import annotations - -import logging -from pathlib import Path -from typing import Any, Optional - -import itk -import pyvista as pv - -from physiomotion4d.register_time_series_images import RegisterTimeSeriesImages -from physiomotion4d.test_tools import TestTools -from physiomotion4d.transform_tools import TransformTools - - -# 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. Wrapping consistently across -# tutorials also matches the style of tutorial_01. -if __name__ == "__main__": - # %% - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - DATA_DIR = REPO_ROOT / "data" - FULL_DATA_DIR = DATA_DIR - TEST_DATA_DIR = DATA_DIR / "test" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_08_dirlab_pca_time_series" - TUTORIAL_07_OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_07_dirlab_pca_model" - CASE: Optional[int] = None - LOG_LEVEL = logging.INFO - - DIRLAB_CASE_PREFIXES = [ - "Case1Pack", - "Case2Pack", - "Case3Pack", - "Case4Pack", - "Case5Pack", - "Case6Pack", - "Case7Pack", - "Case8Deploy", - "Case9Pack", - "Case10Pack", - ] - - def run_tutorial() -> dict[str, Any]: - """Run Tutorial 8: propagate Tutorial 7 PCA fits through DirLab time series. - - Returns - ------- - dict[str, Any] - Per-case transform filenames, propagated mesh filenames, and losses. - """ - - data_dir = TEST_DATA_DIR if TestTools.running_as_test() else FULL_DATA_DIR - output_dir = OUTPUT_DIR - tutorial_07_output_dir = TUTORIAL_07_OUTPUT_DIR - case = CASE - log_level = LOG_LEVEL - - output_dir.mkdir(parents=True, exist_ok=True) - dirlab_dir = data_dir / "DirLab-4DCT" - fits_dir = tutorial_07_output_dir / "fits" - transform_tools = TransformTools(log_level=log_level) - - if case is None: - case_numbers = list(range(1, 11)) - else: - case_numbers = [case] - - tutorial_outputs: dict[str, Any] = {} - for case_number in case_numbers: - case_prefix = DIRLAB_CASE_PREFIXES[case_number - 1] - case_dir = dirlab_dir / f"Case{case_number}" - phase_files = sorted(case_dir.glob("*.mha")) + sorted( - case_dir.glob("*.mhd") - ) - if not phase_files: - phase_files = sorted(dirlab_dir.glob(f"{case_prefix}_T*.mha")) - phase_files += sorted(dirlab_dir.glob(f"{case_prefix}_T*.mhd")) - if not phase_files and case_number == 8: - phase_files = sorted(dirlab_dir.glob("Case8Pack_T*.mha")) - phase_files += sorted(dirlab_dir.glob("Case8Pack_T*.mhd")) - if not phase_files: - print(f"Skipping {case_prefix}: no DirLab phase images found") - continue - - fitted_mesh_file = fits_dir / f"{case_prefix}_lung_lobes_pca_fit.vtp" - if not fitted_mesh_file.exists(): - raise FileNotFoundError( - f"Missing Tutorial 7 fitted mesh: {fitted_mesh_file}. " - "Run Tutorial 7 before Tutorial 8." - ) - - print(f"Registering {case_prefix}: {len(phase_files)} phases") - case_output_dir = output_dir / case_prefix - transforms_dir = case_output_dir / "transforms" - meshes_dir = case_output_dir / "meshes" - transforms_dir.mkdir(parents=True, exist_ok=True) - meshes_dir.mkdir(parents=True, exist_ok=True) - - time_series = [itk.imread(str(phase_file)) for phase_file in phase_files] - fixed_image = time_series[0] - - registrar = RegisterTimeSeriesImages( - registration_method="Greedy_ICON", - log_level=log_level, - ) - registrar.set_modality("ct") - registrar.set_fixed_image(fixed_image) - registration_result = registrar.register_time_series( - moving_images=time_series, - reference_frame=0, - register_reference=False, - ) - - fitted_reference_mesh = pv.read(str(fitted_mesh_file)) - case_transform_files: list[Path] = [] - case_mesh_files: list[Path] = [] - - forward_transforms = registration_result["forward_transforms"] - inverse_transforms = registration_result["inverse_transforms"] - if not ( - len(phase_files) == len(forward_transforms) == len(inverse_transforms) - ): - raise ValueError( - f"{case_prefix}: length mismatch between phase_files " - f"({len(phase_files)}), forward_transforms " - f"({len(forward_transforms)}), and inverse_transforms " - f"({len(inverse_transforms)})." - ) - - for phase_file, phase_to_reference, reference_to_phase in zip( - phase_files, - forward_transforms, - inverse_transforms, - ): - phase_name = phase_file.stem - phase_to_reference_file = transforms_dir / ( - f"{phase_name}_phase_to_reference.hdf" - ) - reference_to_phase_file = transforms_dir / ( - f"{phase_name}_reference_to_phase.hdf" - ) - itk.transformwrite( - phase_to_reference, - str(phase_to_reference_file), - compression=True, - ) - itk.transformwrite( - reference_to_phase, - str(reference_to_phase_file), - compression=True, - ) - - # Warp the reference-space fitted mesh into this phase's space. - # Warping reference -> phase POINTS uses the forward transform - # (the fixed -> moving point map), which is the opposite of the - # transform used to warp an image into phase space (images pull - # back, points push forward). The forward transform is named - # "phase_to_reference" after its image-warp role. See - # docs/developer/transform_conventions. - phase_mesh = transform_tools.transform_pvcontour( - fitted_reference_mesh, - phase_to_reference, - with_deformation_magnitude=True, - ) - phase_mesh_file = meshes_dir / f"{phase_name}_pca_fit.vtp" - phase_mesh.save(phase_mesh_file) - - case_transform_files.extend( - [phase_to_reference_file, reference_to_phase_file] - ) - case_mesh_files.append(phase_mesh_file) - - tutorial_outputs[case_prefix] = { - "reference_phase": phase_files[0], - "phase_files": phase_files, - "fitted_reference_mesh": fitted_mesh_file, - "transform_files": case_transform_files, - "mesh_files": case_mesh_files, - "losses": registration_result["losses"], - } - - return tutorial_outputs - - # %% - # Run this cell in VS Code or Cursor: - tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_09_physicsnemo_mesh_stage_model.py b/tutorials/tutorial_09_physicsnemo_mesh_stage_model.py deleted file mode 100644 index 7e23521..0000000 --- a/tutorials/tutorial_09_physicsnemo_mesh_stage_model.py +++ /dev/null @@ -1,297 +0,0 @@ -""" -Tutorial 9: Train a PhysicsNeMo model for DirLab mesh time-stage prediction. - -This tutorial uses the per-time-point PCA-fitted meshes created by Tutorial 8. -For each case, it trains a small PhysicsNeMo fully connected model that maps -reference mesh point coordinates and a requested normalized respiratory stage to -point displacements. The trained model can then predict a mesh at a new -user-specified stage without running image registration again. - -Data Required -------------- -Run Tutorial 8 first so ``output/tutorial_08_dirlab_pca_time_series`` contains -``Case*/meshes/*_pca_fit.vtp`` files. - -Extra Install Required ----------------------- -PhysicsNeMo is an optional dependency of PhysioMotion4D. Install it with:: - - pip install "physiomotion4d[physicsnemo]" - -PhysicsNeMo itself requires Python >= 3.11. -""" - -# %% -from __future__ import annotations - -import json -import logging -from pathlib import Path -from typing import Any, Optional - -import numpy as np -import pyvista as pv -import torch - - -try: - from physicsnemo.models.mlp import FullyConnected -except ImportError as exc: # pragma: no cover - import-time guard - raise ImportError( - "Tutorial 9 requires PhysicsNeMo, which is an optional dependency. " - 'Install with: pip install "physiomotion4d[physicsnemo]" ' - "(requires Python >= 3.11).", - ) from exc - - -# 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. Wrapping consistently across -# tutorials also matches the style of tutorial_01. -if __name__ == "__main__": - # %% - TUTORIALS_DIR = Path(__file__).resolve().parent - TUTORIAL_08_OUTPUT_DIR = ( - TUTORIALS_DIR / "output" / "tutorial_08_dirlab_pca_time_series" - ) - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_physicsnemo_mesh_stage_model" - TARGET_STAGE = 0.5 - CASE: Optional[int] = None - EPOCHS = 500 - POINTS_PER_MESH = 4096 - LEARNING_RATE = 1.0e-3 - LOG_LEVEL = logging.INFO - - DIRLAB_CASE_PREFIXES = [ - "Case1Pack", - "Case2Pack", - "Case3Pack", - "Case4Pack", - "Case5Pack", - "Case6Pack", - "Case7Pack", - "Case8Deploy", - "Case9Pack", - "Case10Pack", - ] - - def run_tutorial() -> dict[str, Any]: - """Train PhysicsNeMo stage models and predict meshes at ``target_stage``. - - Returns - ------- - dict[str, Any] - Per-case checkpoint, metadata, predicted mesh, and training loss paths. - """ - - tutorial_08_output_dir = TUTORIAL_08_OUTPUT_DIR - output_dir = OUTPUT_DIR - target_stage = TARGET_STAGE - case = CASE - epochs = EPOCHS - points_per_mesh = POINTS_PER_MESH - learning_rate = LEARNING_RATE - log_level = LOG_LEVEL - - logging.basicConfig(level=log_level) - if target_stage < 0.0 or target_stage > 1.0: - raise ValueError("target_stage must be in the normalized range [0.0, 1.0].") - - output_dir.mkdir(parents=True, exist_ok=True) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - num_cases = len(DIRLAB_CASE_PREFIXES) - if case is None: - case_numbers = list(range(1, num_cases + 1)) - else: - if not 1 <= case <= num_cases: - raise ValueError( - f"CASE={case} is out of range; must be an integer between 1 " - f"and {num_cases} (inclusive)." - ) - case_numbers = [case] - - tutorial_outputs: dict[str, Any] = {} - for case_number in case_numbers: - case_prefix = DIRLAB_CASE_PREFIXES[case_number - 1] - mesh_dir = tutorial_08_output_dir / case_prefix / "meshes" - mesh_files = ( - sorted(mesh_dir.glob("*_pca_fit.vtp")) if mesh_dir.exists() else [] - ) - if len(mesh_files) < 2: - message = ( - f"Tutorial 8 output for {case_prefix} is missing or incomplete: " - f"found {len(mesh_files)} '*_pca_fit.vtp' file(s) in {mesh_dir}, " - "expected at least 2. Run Tutorial 8 before Tutorial 9." - ) - if case is not None: - raise FileNotFoundError(message) - logging.info(f"Skipping {case_prefix}: {message}") - continue - - case_output_dir = output_dir / case_prefix - case_output_dir.mkdir(parents=True, exist_ok=True) - - reference_mesh = pv.read(str(mesh_files[0])) - reference_points = np.asarray(reference_mesh.points, dtype=np.float32) - if points_per_mesh <= 0 or points_per_mesh >= reference_mesh.n_points: - point_indices = np.arange(reference_mesh.n_points) - else: - point_indices = np.linspace( - 0, - reference_mesh.n_points - 1, - points_per_mesh, - dtype=np.int64, - ) - - coordinate_mean = reference_points.mean(axis=0) - coordinate_scale = reference_points.std(axis=0) - coordinate_scale = np.where(coordinate_scale == 0.0, 1.0, coordinate_scale) - normalized_reference_points = ( - reference_points[point_indices] - coordinate_mean - ) / coordinate_scale - - training_inputs: list[np.ndarray] = [] - training_targets: list[np.ndarray] = [] - stage_denominator = max(1, len(mesh_files) - 1) - for stage_index, mesh_file in enumerate(mesh_files): - mesh = pv.read(str(mesh_file)) - if mesh.n_points != reference_mesh.n_points: - raise ValueError( - f"{mesh_file} has {mesh.n_points} points, expected " - f"{reference_mesh.n_points}. Tutorial 8 meshes must share topology." - ) - - stage = stage_index / stage_denominator - stage_column = np.full((len(point_indices), 1), stage, dtype=np.float32) - training_inputs.append( - np.hstack([normalized_reference_points, stage_column]) - ) - training_targets.append( - np.asarray(mesh.points[point_indices], dtype=np.float32) - - reference_points[point_indices] - ) - - inputs_array = np.vstack(training_inputs).astype(np.float32) - targets_array = np.vstack(training_targets).astype(np.float32) - displacement_scale = float(np.max(np.abs(targets_array))) - if displacement_scale == 0.0: - displacement_scale = 1.0 - targets_array = targets_array / displacement_scale - - inputs_tensor = torch.from_numpy(inputs_array).to(device) - targets_tensor = torch.from_numpy(targets_array).to(device) - - model = FullyConnected( - in_features=4, - layer_size=128, - out_features=3, - num_layers=4, - activation_fn="silu", - skip_connections=True, - ).to(device) - optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) - loss_function = torch.nn.MSELoss() - - losses: list[float] = [] - model.train() - for epoch in range(epochs): - optimizer.zero_grad() - prediction = model(inputs_tensor) - loss = loss_function(prediction, targets_tensor) - loss.backward() - optimizer.step() - losses.append(float(loss.detach().cpu())) - if epoch == 0 or (epoch + 1) % 100 == 0 or epoch + 1 == epochs: - logging.info( - f"{case_prefix} epoch {epoch + 1:04d}/{epochs}: " - f"loss={losses[-1]:.6f}" - ) - - model.eval() - all_normalized_points = ( - reference_points - coordinate_mean - ) / coordinate_scale - all_stage_column = np.full( - (reference_mesh.n_points, 1), - target_stage, - dtype=np.float32, - ) - prediction_inputs = np.hstack([all_normalized_points, all_stage_column]) - predicted_displacements: list[np.ndarray] = [] - with torch.no_grad(): - for start in range(0, reference_mesh.n_points, 65536): - stop = min(start + 65536, reference_mesh.n_points) - prediction_tensor = torch.from_numpy( - prediction_inputs[start:stop].astype(np.float32) - ).to(device) - predicted_displacements.append( - model(prediction_tensor).cpu().numpy() * displacement_scale - ) - - predicted_mesh = reference_mesh.copy(deep=True) - predicted_mesh.points = reference_points + np.vstack( - predicted_displacements - ) - - stage_tag = f"{target_stage:.3f}".replace(".", "p") - checkpoint_file = case_output_dir / "physicsnemo_stage_model.pt" - metadata_file = case_output_dir / "physicsnemo_stage_model_metadata.json" - losses_file = case_output_dir / "training_losses.json" - predicted_mesh_file = ( - case_output_dir / f"{case_prefix}_stage_{stage_tag}.vtp" - ) - - torch.save( - { - "model_state_dict": model.state_dict(), - "coordinate_mean": coordinate_mean.tolist(), - "coordinate_scale": coordinate_scale.tolist(), - "displacement_scale": displacement_scale, - "target_stage": target_stage, - "mesh_files": [str(mesh_file) for mesh_file in mesh_files], - }, - checkpoint_file, - ) - metadata_file.write_text( - json.dumps( - { - "architecture": "physicsnemo.models.mlp.FullyConnected", - "input_features": [ - "reference_x_normalized", - "reference_y_normalized", - "reference_z_normalized", - "normalized_stage", - ], - "output_features": ["dx", "dy", "dz"], - "target_stage": target_stage, - "epochs": epochs, - "points_per_mesh": len(point_indices), - "learning_rate": learning_rate, - "coordinate_mean": coordinate_mean.tolist(), - "coordinate_scale": coordinate_scale.tolist(), - "displacement_scale": displacement_scale, - "training_meshes": [str(mesh_file) for mesh_file in mesh_files], - }, - indent=2, - ), - encoding="utf-8", - ) - losses_file.write_text(json.dumps(losses, indent=2), encoding="utf-8") - predicted_mesh.save(predicted_mesh_file) - - tutorial_outputs[case_prefix] = { - "checkpoint_file": checkpoint_file, - "metadata_file": metadata_file, - "losses_file": losses_file, - "predicted_mesh_file": predicted_mesh_file, - "final_loss": losses[-1], - } - - return tutorial_outputs - - # %% - # Run this cell in VS Code or Cursor: - tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py b/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py new file mode 100644 index 0000000..9885c5c --- /dev/null +++ b/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py @@ -0,0 +1,878 @@ +""" +Tutorial 9a (MGN): MeshGraphNet model for cardiac mesh stage prediction. + +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +Drop-in companion to the MLP trainer Tutorial 9b +(``tutorial_09b_cardiac_train_physicsnemo_mlp.py``). It uses the same +per-time-point SSM surfaces created by Tutorial 8 +(``tutorial_08_cardiac_fit_model.py``) and the same Option B displacement +convention, but replaces the FullyConnected MLP with a PhysicsNeMo MeshGraphNet +that explicitly exploits the surface mesh topology via message passing between +neighbouring vertices. Outputs land in ``OUTPUT_DIR`` (``output_mgn/``) so the +two models can be evaluated and compared side by side. Evaluate the trained +model with Tutorial 10a (``tutorial_10a_cardiac_eval_physicsnemo_mgn.py``). + +Why a GNN? +---------- +The SSM mesh has a fixed, consistent topology across all subjects. Cardiac tissue is a +physical continuum - adjacent vertices co-vary smoothly. The MLP trainer must infer +this from xyz coordinates alone. MeshGraphNet encodes the prior directly by passing +messages along mesh edges, giving the model an explicit continuum-deformation inductive +bias. + +Node features (per vertex): [norm_x, norm_y, norm_z, pca_c1 ... pca_cN, stage] +Edge features (per edge): [rel_x, rel_y, rel_z, distance] (derived from mean shape) +Output (per vertex): [dx, dy, dz] (displacement in mm, after rescaling) + +The edge topology is extracted once from the mean-shape surface and shared across every +(subject, phase) sample. PyTorch Geometric's ``Batch.from_data_list`` handles +mini-batching by joining disconnected sub-graphs. The shared graph topology +(``shared_edge_index.pt`` / ``shared_edge_features.pt``) is saved alongside the weights +so Tutorial 10a (``tutorial_10a_cardiac_eval_physicsnemo_mgn.py``) can replay it +at inference time. + +Bring Your Own Data +------------------- +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioMotion4D/`` layout produced by Tutorial 8, not at the repository +``data/`` directory. Edit them to match your own data location. + +Data Required +------------- +Run Tutorial 8 (``tutorial_08_cardiac_fit_model.py``) first (same requirement as +the MLP trainer). + +Extra Install Required +---------------------- +PhysicsNeMo and PyTorch Geometric must be installed:: + + pip install "physiomotion4d[physicsnemo]" + pip install torch-geometric + +``torch_scatter`` (a PyTorch Geometric backend) must be built from source when using a +custom NVIDIA PyTorch build because no matching pre-built wheel exists on data.pyg.org:: + + pip install torch-scatter --no-build-isolation +""" + +# %% +from __future__ import annotations + +import csv +import json +import logging +import shutil +import sys +from collections import defaultdict +from pathlib import Path +from typing import Any, Optional, cast + +import numpy as np +import pyvista as pv +import torch + +from physiomotion4d.test_tools import TestTools + +try: + import torch_geometric.utils as pyg_utils + from torch_geometric.data import Batch, Data + + from physicsnemo.models.meshgraphnet import MeshGraphNet +except ImportError as exc: + raise ImportError( + "Tutorial 9a requires PhysicsNeMo and PyTorch Geometric. Install with:\n" + ' pip install "physiomotion4d[physicsnemo]"\n' + " pip install torch-geometric" + ) from exc + + +if __name__ == "__main__": + # %% + TUTORIALS_DIR = Path(__file__).resolve().parent + FITTED_MESHES_DIR = Path("D:/PhysioMotion4D/duke_data/fitted_kcl_meshes") + PCA_MEAN_VTU = Path("D:/PhysioMotion4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + EPOCHS = 1500 + OUTPUT_DIR = TUTORIALS_DIR / "output_mgn" + RMSE_LOG_INTERVAL = 100 + LOSS_LOG_INTERVAL = ( + 1 # print every epoch so we can measure per-epoch time immediately + ) + # Mini-batch size in (subject, phase) *graphs*. + # concat_efeat stores (BxE, 3H) FP32 per processor step x PROCESSOR_SIZE steps: + # PROCESSOR_SIZE=3, B=4, H=128: 3 x 4M x 384 x 4 = 18.4 GB -> safe, good GPU util + # PROCESSOR_SIZE=10, B=2, H=128: 10 x 2M x 384 x 4 = 30.7 GB -> safe but slow (83 h) + # 3 hops is sufficient for local mesh-continuity; 10 hops adds marginal benefit here. + BATCH_SIZE_GRAPHS = 4 + LEARNING_RATE = 1.0e-3 + # MeshGraphNet hyper-parameters + PROCESSOR_SIZE = 3 # 3 message-passing hops: enough for surface continuity, + # ~3x faster than 10; 10K epochs ~ 8 h vs 83 h + HIDDEN_DIM = 128 # sufficient capacity for 27 training subjects + # Gradient checkpointing (0 = disabled). + # checkpointing=5 caused the training loop to stall: Batch.from_data_list was called + # 380K times (38 batches x 10K epochs), each round-tripping 8M edges CPU<->GPU. + # With B=2 the full activation storage fits without checkpointing. + NUM_PROCESSOR_CHECKPOINT_SEGMENTS = 0 + NUM_LAYERS_PROCESSOR = 2 # MLP layers inside each processor step + NUM_LAYERS_ENCODER = 2 + NUM_LAYERS_DECODER = 2 + + TEST_SUBJECTS: Optional[list[str]] = ["pm0028"] + VAL_SUBJECTS: Optional[list[str]] = ["pm0027"] + USE_MEAN_SHAPE_COORDS = True + LOG_LEVEL = logging.INFO + RESUME_FROM_WEIGHTS: Optional[Path] = None + + # ---------------------------------------------------------------------- # + + def _next_output_dir(base: Path) -> Path: + if not base.exists(): + return base + n = 1 + while True: + candidate = base.parent / f"{base.name}_{n}" + if not candidate.exists(): + return candidate + n += 1 + + def _gating_stage_from_filename(mesh_file: Path) -> float: + stem = mesh_file.stem + for part in stem.split("_"): + if part.startswith("g") and part[1:].isdigit(): + return int(part[1:]) / 100.0 + raise ValueError(f"Cannot parse gating percentage from filename: {mesh_file}") + + def _uncompiled_state_dict(model: torch.nn.Module) -> dict: + """Return the base model's state dict, unwrapping torch.compile if applied.""" + return cast(dict, getattr(model, "_orig_mod", model).state_dict()) + + def _mesh_to_edge_index(poly: pv.PolyData) -> torch.Tensor: + """Extract undirected edge_index from triangulated PyVista PolyData faces.""" + faces = poly.faces.reshape(-1, 4)[:, 1:] # (F, 3) - strip leading count + src = np.concatenate([faces[:, 0], faces[:, 1], faces[:, 2]]) + dst = np.concatenate([faces[:, 1], faces[:, 2], faces[:, 0]]) + edge_index = torch.tensor(np.stack([src, dst]), dtype=torch.long) + return cast(torch.Tensor, pyg_utils.to_undirected(edge_index)) + + def _compute_edge_features( + coords: np.ndarray, edge_index: torch.Tensor + ) -> torch.Tensor: + """Build (N_edges, 4) edge feature tensor: [rel_x, rel_y, rel_z, distance].""" + ei = edge_index.numpy() + disp = coords[ei[1]] - coords[ei[0]] # (N_edges, 3) + dist = np.linalg.norm(disp, axis=1, keepdims=True) # (N_edges, 1) + return torch.tensor(np.hstack([disp, dist]), dtype=torch.float32) + + def _batched_rmse_mm( + model: MeshGraphNet, + node_feats_gpu: torch.Tensor, # (N_samples, n_mesh_points, in_features) + targets_gpu: torch.Tensor, # (N_samples, n_mesh_points, 3) + full_batch_graph: "Data", + full_edge_feats: torch.Tensor, + partial_batch_graph: "Data", + partial_edge_feats: torch.Tensor, + displacement_scale: float, + batch_size: int, + n_mesh_points: int, + in_features: int, + ) -> float: + """Euclidean RMSE in mm over pre-stacked GPU tensors (no CPU transfers).""" + n_samples = node_feats_gpu.shape[0] + total_sq = 0.0 + n_total = 0 + with torch.no_grad(): + for start in range(0, n_samples, batch_size): + end = min(start + batch_size, n_samples) + b = end - start + nf = node_feats_gpu[start:end].reshape(b * n_mesh_points, in_features) + tgt = targets_gpu[start:end].reshape(b * n_mesh_points, 3) + bg = full_batch_graph if b == batch_size else partial_batch_graph + ef = full_edge_feats if b == batch_size else partial_edge_feats + pred = model(nf, ef, bg) + err_mm = (pred - tgt) * displacement_scale + total_sq += float(torch.sum(err_mm**2)) + n_total += b * n_mesh_points + return float(np.sqrt(total_sq / n_total)) + + def _infer_all_points( + model: MeshGraphNet, + norm_coords: np.ndarray, + norm_pca: np.ndarray, + stage: float, + shared_graph: Data, + shared_edge_feats: torch.Tensor, + displacement_scale: float, + device: torch.device, + ) -> np.ndarray: + """Run inference for a single (subject, phase) sample; return displacements (mm).""" + n = len(norm_coords) + pca_tile = np.tile(norm_pca, (n, 1)) + stage_col = np.full((n, 1), stage, dtype=np.float32) + node_feats = torch.tensor( + np.hstack([norm_coords, pca_tile, stage_col]), dtype=torch.float32 + ).to(device) + graph = shared_graph.clone().to(device) + edge_feats = shared_edge_feats.to(device) + with torch.no_grad(): + pred = model(node_feats, edge_feats, graph) + return np.asarray(pred.cpu().numpy()) * displacement_scale + + def run_tutorial() -> dict[str, Any]: + """Train a MeshGraphNet model across all subjects and evaluate. + + Same input/output convention as the MLP trainer (Option B displacements + relative to each subject's SSM reference surface) so results are directly + comparable. + """ + fitted_meshes_dir = FITTED_MESHES_DIR + pca_mean_vtu = PCA_MEAN_VTU + epochs = EPOCHS + rmse_log_interval = RMSE_LOG_INTERVAL + loss_log_interval = LOSS_LOG_INTERVAL + batch_size_graphs = BATCH_SIZE_GRAPHS + learning_rate = LEARNING_RATE + test_subjects = TEST_SUBJECTS + val_subjects = VAL_SUBJECTS + use_mean_shape_coords = USE_MEAN_SHAPE_COORDS + resume_from_weights = RESUME_FROM_WEIGHTS + + logging.basicConfig(level=LOG_LEVEL) + + # In test mode, train for only a couple of epochs so the tutorial test + # exercises the full pipeline without a long GPU run. + if TestTools.running_as_test(): + epochs = min(epochs, 2) + logging.info("Test mode: reducing epochs to %d", epochs) + + if resume_from_weights is not None: + output_dir = _next_output_dir(OUTPUT_DIR) + logging.info(f"Resuming from {resume_from_weights}; output: {output_dir}") + else: + output_dir = OUTPUT_DIR + test_output_dir = output_dir / "test_predictions" + output_dir.mkdir(parents=True, exist_ok=True) + test_output_dir.mkdir(parents=True, exist_ok=True) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # ------------------------------------------------------------------ # + # 1. Load PCA mean-shape surface (shared coordinate and topology) # + # ------------------------------------------------------------------ # + logging.info(f"Loading PCA mean shape from {pca_mean_vtu}") + mean_vol = pv.read(str(pca_mean_vtu)) + mean_surf: pv.PolyData = mean_vol.extract_surface(algorithm="dataset_surface") + mean_shape_coords = np.asarray(mean_surf.points, dtype=np.float32) + logging.info( + f"Mean shape surface: {len(mean_shape_coords)} points, " + f"{mean_surf.n_faces} faces" + ) + + # ------------------------------------------------------------------ # + # 2. Build shared graph topology and edge features (from mean shape) # + # ------------------------------------------------------------------ # + logging.info("Building shared mesh graph from mean-shape faces ...") + shared_edge_index = _mesh_to_edge_index(mean_surf) + shared_edge_feats = _compute_edge_features(mean_shape_coords, shared_edge_index) + n_mesh_edges = shared_edge_index.shape[1] + logging.info( + f"Graph: {len(mean_shape_coords)} nodes, {n_mesh_edges} edges " + f"(~{n_mesh_edges / mean_surf.n_faces:.1f}x n_faces - expected ~6 for triangles)" + ) + # shared_graph holds connectivity only; node/edge features passed separately + shared_graph = Data( + edge_index=shared_edge_index, + num_nodes=len(mean_shape_coords), + ) + + # ------------------------------------------------------------------ # + # 3. Discover and load all subjects # + # ------------------------------------------------------------------ # + subjects: dict[str, dict] = {} + for subject_dir in sorted(fitted_meshes_dir.glob("pm????")): + sid = subject_dir.name + ref_file = subject_dir / f"{sid}_ssm_surface.vtp" + pca_file = subject_dir / f"{sid}_ssm_pca_coefficients.json" + mesh_files = sorted(subject_dir.glob(f"{sid}_g0*_ssm_surface.vtp")) + + missing = [p for p in (ref_file, pca_file) if not p.exists()] + if missing or len(mesh_files) < 2: + msg = ( + f"Skipping {sid}: missing files {[str(p) for p in missing]}" + if missing + else f"Skipping {sid}: only {len(mesh_files)} gated phase(s) found" + ) + logging.info(msg) + continue + + ref_mesh = pv.read(str(ref_file)) + subjects[sid] = { + "subject_dir": subject_dir, + "ref_mesh": ref_mesh, + "ref_points": np.asarray(ref_mesh.points, dtype=np.float32), + "pca_coeffs": np.array( + json.loads(pca_file.read_text(encoding="utf-8")), dtype=np.float32 + ), + "mesh_files": mesh_files, + } + + if len(subjects) < 3: + raise RuntimeError( + f"Found only {len(subjects)} valid subject(s); need at least 3." + ) + + n_pca = next(iter(subjects.values()))["pca_coeffs"].shape[0] + n_mesh_points = next(iter(subjects.values()))["ref_points"].shape[0] + + if n_mesh_points != len(mean_shape_coords): + raise RuntimeError( + f"SSM surfaces have {n_mesh_points} points but mean shape has " + f"{len(mean_shape_coords)} - topology mismatch." + ) + + # ------------------------------------------------------------------ # + # 4. Validate splits # + # ------------------------------------------------------------------ # + all_sids = set(subjects.keys()) + test_list = test_subjects if test_subjects is not None else [] + val_list = val_subjects if val_subjects is not None else [] + + unknown = [s for s in test_list + val_list if s not in all_sids] + if unknown: + raise ValueError( + f"Split subjects not found in {fitted_meshes_dir}: {unknown}" + ) + overlap = set(test_list) & set(val_list) + if overlap: + raise ValueError(f"Subjects in both TEST and VAL splits: {sorted(overlap)}") + + test_sids: set[str] = set(test_list) + val_sids: set[str] = set(val_list) + train_sids: set[str] = all_sids - test_sids - val_sids + if not train_sids: + raise ValueError("No subjects remain for training.") + + logging.info( + f"Subject split - train: {len(train_sids)}, " + f"val: {len(val_sids)}, test: {len(test_sids)}" + ) + + # ------------------------------------------------------------------ # + # 5. Normalisation statistics # + # ------------------------------------------------------------------ # + resume_ckpt: Optional[dict] = None + if resume_from_weights is not None: + logging.info(f"Loading prior weights from {resume_from_weights}") + resume_ckpt = torch.load( + str(resume_from_weights), map_location="cpu", weights_only=True + ) + coordinate_mean = np.array(resume_ckpt["coordinate_mean"], dtype=np.float32) + coordinate_scale = np.array( + resume_ckpt["coordinate_scale"], dtype=np.float32 + ) + pca_mean_vec = np.array(resume_ckpt["pca_mean"], dtype=np.float32) + pca_scale_vec = np.array(resume_ckpt["pca_scale"], dtype=np.float32) + displacement_scale = float(resume_ckpt["displacement_scale"]) + else: + coord_ref = ( + mean_shape_coords + if use_mean_shape_coords + else np.vstack([subjects[s]["ref_points"] for s in train_sids]) + ) + coordinate_mean = coord_ref.mean(axis=0) + coordinate_scale = coord_ref.std(axis=0) + coordinate_scale = np.where(coordinate_scale == 0.0, 1.0, coordinate_scale) + + train_pca = np.vstack([subjects[s]["pca_coeffs"] for s in train_sids]) + pca_mean_vec = train_pca.mean(axis=0) + pca_scale_vec = train_pca.std(axis=0) + pca_scale_vec = np.where(pca_scale_vec == 0.0, 1.0, pca_scale_vec) + + # ------------------------------------------------------------------ # + # 6. Build sample lists: one entry per (subject, phase) graph # + # ------------------------------------------------------------------ # + logging.info("Building per-sample graph feature lists ...") + + def _build_samples( + sids: set[str], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Return list of (node_feats, targets) tensors for each (subject, phase).""" + samples: list[tuple[torch.Tensor, torch.Tensor]] = [] + for sid in sorted(sids): + data = subjects[sid] + coords = ( + mean_shape_coords if use_mean_shape_coords else data["ref_points"] + ) + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (data["pca_coeffs"] - pca_mean_vec) / pca_scale_vec + pca_tile = np.tile(norm_pca, (n_mesh_points, 1)) + for mesh_file in data["mesh_files"]: + mesh = pv.read(str(mesh_file)) + if mesh.n_points != n_mesh_points: + raise ValueError( + f"{mesh_file} has {mesh.n_points} points, " + f"expected {n_mesh_points}." + ) + stage = _gating_stage_from_filename(mesh_file) + stage_col = np.full((n_mesh_points, 1), stage, dtype=np.float32) + node_feats = torch.tensor( + np.hstack([norm_coords, pca_tile, stage_col]), + dtype=torch.float32, + ) + targets_raw = ( + np.asarray(mesh.points, dtype=np.float32) - data["ref_points"] + ) + samples.append( + (node_feats, torch.tensor(targets_raw, dtype=torch.float32)) + ) + return samples + + train_samples = _build_samples(train_sids) + val_samples = _build_samples(val_sids) + + # Derive displacement_scale from training targets (or inherit from checkpoint). + if resume_ckpt is None: + all_targets = torch.cat([s[1] for s in train_samples]) + displacement_scale = float(torch.max(torch.abs(all_targets))) + if displacement_scale == 0.0: + displacement_scale = 1.0 + + # Normalise targets in-place. + train_samples = [(n, t / displacement_scale) for n, t in train_samples] + val_samples = [(n, t / displacement_scale) for n, t in val_samples] + + in_features = 3 + n_pca + 1 + logging.info( + f"Training samples: {len(train_samples)}, val: {len(val_samples)}, " + f"in_features={in_features}, displacement_scale={displacement_scale:.4f} mm" + ) + + # ------------------------------------------------------------------ # + # 7. Model # + # ------------------------------------------------------------------ # + model = MeshGraphNet( + input_dim_nodes=in_features, + input_dim_edges=4, # rel_x, rel_y, rel_z, distance + output_dim=3, + processor_size=PROCESSOR_SIZE, + hidden_dim_processor=HIDDEN_DIM, + hidden_dim_node_encoder=HIDDEN_DIM, + num_layers_node_encoder=NUM_LAYERS_ENCODER, + hidden_dim_node_decoder=HIDDEN_DIM, + num_layers_node_decoder=NUM_LAYERS_DECODER, + hidden_dim_edge_encoder=HIDDEN_DIM, + num_layers_edge_processor=NUM_LAYERS_PROCESSOR, + num_layers_node_processor=NUM_LAYERS_PROCESSOR, + aggregation="mean", + num_processor_checkpoint_segments=NUM_PROCESSOR_CHECKPOINT_SEGMENTS, + ).to(device) + + if resume_ckpt is not None: + state = resume_ckpt.get("model_state_dict", resume_ckpt) + model.load_state_dict(state) + logging.info("Loaded model weights from prior checkpoint.") + + # torch.compile: Linux-only (Triton unavailable on Windows). + if sys.platform != "win32": + try: + model = torch.compile(model) + logging.info("torch.compile enabled.") + except Exception as _e: + logging.info(f"torch.compile skipped ({_e}).") + else: + logging.info("torch.compile skipped on Windows.") + + # Move shared graph tensors to GPU once (used for single-sample inference). + shared_edge_feats_gpu = shared_edge_feats.to(device) + shared_graph_gpu = shared_graph.clone().to(device) + + # Pre-stack all training/val tensors onto GPU (~5 GB total - fits easily). + # Eliminates per-batch CPU->GPU transfers that were a primary stall source. + logging.info("Pre-stacking training and validation tensors onto GPU ...") + train_node_feats_gpu = torch.stack([s[0] for s in train_samples]).to(device) + train_targets_gpu = torch.stack([s[1] for s in train_samples]).to(device) + n_val = len(val_samples) + has_val = n_val > 0 + if has_val: + val_node_feats_gpu = torch.stack([s[0] for s in val_samples]).to(device) + val_targets_gpu = torch.stack([s[1] for s in val_samples]).to(device) + else: + logging.info("No validation subjects configured; skipping val RMSE.") + n_train = len(train_samples) + + # Pre-build batched graph and edge features for the full-batch size (and for any + # partial last batch). All samples share the same mesh topology, so these can be + # built once and reused every step of every epoch - the previous code rebuilt them + # ~380 K times (38 batches x 10 K epochs), each involving a costly CPU/GPU round + # trip across 8 M edge indices. + logging.info( + "Pre-building batched graph and edge features (built once, reused every step) ..." + ) + full_batch_graph = Batch.from_data_list([shared_graph] * batch_size_graphs).to( + device + ) + full_edge_feats = shared_edge_feats.repeat(batch_size_graphs, 1).to(device) + n_partial = n_train % batch_size_graphs + if n_partial > 0: + partial_batch_graph = Batch.from_data_list([shared_graph] * n_partial).to( + device + ) + partial_edge_feats = shared_edge_feats.repeat(n_partial, 1).to(device) + else: + partial_batch_graph = full_batch_graph + partial_edge_feats = full_edge_feats + if has_val: + n_val_partial = n_val % batch_size_graphs + if n_val_partial > 0: + val_partial_batch_graph = Batch.from_data_list( + [shared_graph] * n_val_partial + ).to(device) + val_partial_edge_feats = shared_edge_feats.repeat(n_val_partial, 1).to( + device + ) + else: + val_partial_batch_graph = full_batch_graph + val_partial_edge_feats = full_edge_feats + + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + loss_fn = torch.nn.MSELoss() + + losses: list[float] = [] + rmse_log: list[dict] = [] + + import time as _time + + # ------------------------------------------------------------------ # + # 8. Training loop # + # ------------------------------------------------------------------ # + for epoch in range(epochs): + _t0 = _time.perf_counter() + model.train() + epoch_loss = torch.zeros( + (), device=device + ) # accumulate on GPU, one sync/epoch + # Shuffle indices on GPU - no CPU involvement. + perm = torch.randperm(n_train, device=device) + + for start in range(0, n_train, batch_size_graphs): + idx = perm[start : start + batch_size_graphs] + b = int(idx.shape[0]) + + # Index pre-stacked GPU tensors; reshape is a view (no copy). + node_feats = train_node_feats_gpu[idx].reshape( + b * n_mesh_points, in_features + ) + targets = train_targets_gpu[idx].reshape(b * n_mesh_points, 3) + batch_graph = ( + full_batch_graph if b == batch_size_graphs else partial_batch_graph + ) + edge_feats = ( + full_edge_feats if b == batch_size_graphs else partial_edge_feats + ) + + optimizer.zero_grad(set_to_none=True) + with torch.amp.autocast(device.type, dtype=torch.bfloat16): + pred = model(node_feats, edge_feats, batch_graph) + loss = loss_fn(pred, targets) + loss.backward() + optimizer.step() + epoch_loss = epoch_loss + loss.detach() * (b * n_mesh_points) + + losses.append( + float(epoch_loss / (n_train * n_mesh_points)) + ) # one GPU sync/epoch + + _epoch_s = _time.perf_counter() - _t0 + if (epoch + 1) % loss_log_interval == 0: + logging.info( + " epoch %05d/%d loss=%.6f %.1fs/epoch ETA %.1fh", + epoch + 1, + epochs, + losses[-1], + _epoch_s, + _epoch_s * (epochs - epoch - 1) / 3600, + ) + + if (epoch + 1) % rmse_log_interval == 0 or epoch + 1 == epochs: + model.eval() + train_rmse = _batched_rmse_mm( + model, + train_node_feats_gpu, + train_targets_gpu, + full_batch_graph, + full_edge_feats, + partial_batch_graph, + partial_edge_feats, + displacement_scale, + batch_size_graphs, + n_mesh_points, + in_features, + ) + val_rmse = ( + _batched_rmse_mm( + model, + val_node_feats_gpu, + val_targets_gpu, + full_batch_graph, + full_edge_feats, + val_partial_batch_graph, + val_partial_edge_feats, + displacement_scale, + batch_size_graphs, + n_mesh_points, + in_features, + ) + if has_val + else float("nan") + ) + rmse_log.append( + { + "epoch": epoch + 1, + "train_rmse_mm": train_rmse, + "val_rmse_mm": val_rmse, + } + ) + ckpt_path = output_dir / f"mgn_stage_model_epoch_{epoch + 1:05d}.pt" + torch.save(_uncompiled_state_dict(model), ckpt_path) + logging.info( + "INTERMITTENT TEST epoch %05d/%d " + "train RMSE=%.4f mm val RMSE=%.4f mm checkpoint=%s", + epoch + 1, + epochs, + train_rmse, + val_rmse, + ckpt_path.name, + ) + model.eval() + + # ------------------------------------------------------------------ # + # 9. Save model weights + metadata # + # ------------------------------------------------------------------ # + checkpoint_file = output_dir / "mgn_stage_model.pt" + metadata_file = output_dir / "mgn_stage_model_metadata.json" + + torch.save( + { + "model_state_dict": _uncompiled_state_dict(model), + "in_features": in_features, + "processor_size": PROCESSOR_SIZE, + "hidden_dim": HIDDEN_DIM, + "coordinate_mean": coordinate_mean.tolist(), + "coordinate_scale": coordinate_scale.tolist(), + "pca_mean": pca_mean_vec.tolist(), + "pca_scale": pca_scale_vec.tolist(), + "displacement_scale": displacement_scale, + "use_mean_shape_coords": use_mean_shape_coords, + "train_subject_ids": sorted(train_sids), + "val_subject_ids": sorted(val_sids), + "test_subject_ids": sorted(test_sids), + "resumed_from": str(resume_from_weights) + if resume_from_weights + else None, + }, + checkpoint_file, + ) + # Save shared graph topology for inference replay. + torch.save(shared_edge_index, output_dir / "shared_edge_index.pt") + torch.save(shared_edge_feats, output_dir / "shared_edge_features.pt") + + input_feature_names = ( + ( + ["mean_shape_x", "mean_shape_y", "mean_shape_z"] + if use_mean_shape_coords + else ["ssm_x", "ssm_y", "ssm_z"] + ) + + [f"pca_c{i + 1}" for i in range(n_pca)] + + ["stage"] + ) + metadata_file.write_text( + json.dumps( + { + "architecture": "physicsnemo.models.meshgraphnet.MeshGraphNet", + "input_node_features": input_feature_names, + "input_edge_features": ["rel_x", "rel_y", "rel_z", "distance"], + "output_features": ["dx", "dy", "dz"], + "n_subjects": len(subjects), + "n_mesh_points": n_mesh_points, + "n_mesh_edges": n_mesh_edges, + "in_features": in_features, + "processor_size": PROCESSOR_SIZE, + "hidden_dim": HIDDEN_DIM, + "epochs": epochs, + "learning_rate": learning_rate, + "coordinate_mean": coordinate_mean.tolist(), + "coordinate_scale": coordinate_scale.tolist(), + "pca_mean": pca_mean_vec.tolist(), + "pca_scale": pca_scale_vec.tolist(), + "displacement_scale": displacement_scale, + "use_mean_shape_coords": use_mean_shape_coords, + "displacement_convention": "Option B: relative to subject ssm_surface", + "resumed_from": str(resume_from_weights) + if resume_from_weights + else None, + }, + indent=2, + ), + encoding="utf-8", + ) + (output_dir / "training_losses.json").write_text( + json.dumps(losses, indent=2), encoding="utf-8" + ) + rmse_log_file = output_dir / "training_validation_rmse.json" + rmse_log_file.write_text(json.dumps(rmse_log, indent=2), encoding="utf-8") + rmse_csv_file = output_dir / "training_validation_rmse.csv" + with rmse_csv_file.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter( + fh, fieldnames=["epoch", "train_rmse_mm", "val_rmse_mm"] + ) + writer.writeheader() + writer.writerows(rmse_log) + + # Copy PCA assets so the output directory is self-contained. + pca_src_dir = pca_mean_vtu.parent + shutil.copy2(pca_mean_vtu, output_dir / pca_mean_vtu.name) + pca_model_src = pca_src_dir / "pca_model.json" + if pca_model_src.exists(): + shutil.copy2(pca_model_src, output_dir / pca_model_src.name) + mean_surf.save(str(output_dir / "pca_mean_surface.vtp")) + + logging.info(f"Model saved to {checkpoint_file}") + + # ------------------------------------------------------------------ # + # 10. Evaluate test and val subjects # + # ------------------------------------------------------------------ # + def _evaluate_subject(sid: str, split_label: str) -> tuple[list[dict], Path]: + data = subjects[sid] + coords = mean_shape_coords if use_mean_shape_coords else data["ref_points"] + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (data["pca_coeffs"] - pca_mean_vec) / pca_scale_vec + + subj_out_dir = output_dir / sid + subj_out_dir.mkdir(parents=True, exist_ok=True) + + sq_err_sum = np.zeros(n_mesh_points, dtype=np.float64) + stats = [] + + for phase_file in data["mesh_files"]: + stage = _gating_stage_from_filename(phase_file) + pred_disps = _infer_all_points( + model, + norm_coords, + norm_pca, + stage, + shared_graph_gpu, + shared_edge_feats_gpu, + displacement_scale, + device, + ) + pred_points = data["ref_points"] + pred_disps + + pred_mesh = data["ref_mesh"].copy(deep=True) + pred_mesh.points = pred_points + gating_tag = phase_file.stem.split("_ssm_surface")[0].split("_")[-1] + pred_mesh.save( + subj_out_dir / f"{sid}_{gating_tag}_ssm_surface_pred.vtp" + ) + + actual_points = np.asarray( + pv.read(str(phase_file)).points, dtype=np.float32 + ) + errors = pred_points - actual_points + euclidean = np.linalg.norm(errors, axis=1) + sq_err_sum += euclidean.astype(np.float64) ** 2 + stats.append( + { + "subject_id": sid, + "split": split_label, + "gating_tag": gating_tag, + "stage": stage, + "n_points": len(euclidean), + "mean_error_mm": float(euclidean.mean()), + "median_error_mm": float(np.median(euclidean)), + "max_error_mm": float(euclidean.max()), + "rms_error_mm": float(np.sqrt(np.mean(euclidean**2))), + "std_error_mm": float(euclidean.std()), + "mean_abs_error_x_mm": float(np.abs(errors[:, 0]).mean()), + "mean_abs_error_y_mm": float(np.abs(errors[:, 1]).mean()), + "mean_abs_error_z_mm": float(np.abs(errors[:, 2]).mean()), + } + ) + logging.info( + f"{sid} [{split_label}] {gating_tag}: " + f"mean={stats[-1]['mean_error_mm']:.3f} mm " + f"max={stats[-1]['max_error_mm']:.3f} mm" + ) + + point_rmse = np.sqrt(sq_err_sum / len(data["mesh_files"])).astype( + np.float32 + ) + rmse_mesh = data["ref_mesh"].copy(deep=True) + rmse_mesh.point_data["RMSE_mm"] = point_rmse + rmse_file = subj_out_dir / f"{sid}_ssm_surface_rmse.vtp" + rmse_mesh.save(rmse_file) + logging.info( + f"{sid} [{split_label}] per-point RMSE: " + f"mean={point_rmse.mean():.3f} mm max={point_rmse.max():.3f} mm" + ) + return stats, rmse_file + + all_stats = [] + tutorial_outputs = {} + + for sid in sorted(test_sids): + stats, rmse_file = _evaluate_subject(sid, "test") + all_stats.extend(stats) + tutorial_outputs[sid] = { + "split": "test", + "rmse_file": rmse_file, + "final_loss": losses[-1], + "n_phases": len(subjects[sid]["mesh_files"]), + } + + for sid in sorted(val_sids): + stats, rmse_file = _evaluate_subject(sid, "val") + all_stats.extend(stats) + tutorial_outputs[sid] = { + "split": "val", + "rmse_file": rmse_file, + "final_loss": losses[-1], + "n_phases": len(subjects[sid]["mesh_files"]), + } + + if all_stats: + per_phase_csv = test_output_dir / "statistics_per_phase.csv" + with per_phase_csv.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(all_stats[0].keys())) + writer.writeheader() + writer.writerows(all_stats) + + subject_rows = defaultdict(list) + for row in all_stats: + subject_rows[row["subject_id"]].append(row) + + summary_rows = [ + { + "subject_id": sid, + "split": rows[0]["split"], + "n_phases": len(rows), + "mean_error_mm": float(np.mean([r["mean_error_mm"] for r in rows])), + "mean_max_error_mm": float( + np.mean([r["max_error_mm"] for r in rows]) + ), + "overall_max_error_mm": float( + np.max([r["max_error_mm"] for r in rows]) + ), + "mean_rms_error_mm": float( + np.mean([r["rms_error_mm"] for r in rows]) + ), + } + for sid, rows in sorted(subject_rows.items()) + ] + summary_csv = test_output_dir / "statistics_summary.csv" + with summary_csv.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(summary_rows[0].keys())) + writer.writeheader() + writer.writerows(summary_rows) + + return tutorial_outputs + + # %% + tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py b/tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py new file mode 100644 index 0000000..2e3c366 --- /dev/null +++ b/tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py @@ -0,0 +1,805 @@ +""" +Tutorial 9b (MLP): Train a PhysicsNeMo MLP for cardiac mesh stage prediction. + +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +This tutorial consumes the per-time-point SSM-warped surfaces created by +Tutorial 8 (``tutorial_08_cardiac_fit_model.py``). It trains a single shared +PhysicsNeMo fully connected (MLP) model across all training subjects that maps a +surface point ``(x, y, z, pca_c1 ... pca_cN, stage)`` to the point's +displacement from that subject's SSM reference surface, where ``stage`` is the +normalized cardiac stage (RR-interval fraction). Once trained, the model +predicts a cardiac mesh at any requested stage without re-running image +registration. Evaluate the trained model with Tutorial 10b +(``tutorial_10b_cardiac_eval_physicsnemo_mlp.py``). + +The companion Tutorial 9a (``tutorial_09a_cardiac_train_physicsnemo_mgn.py``) +solves the same task with a MeshGraphNet so the two architectures can be compared +directly; both use the same Option B displacement convention (targets relative to +each subject's own SSM reference surface). Subjects are split into +train / val / test via the explicit ``TEST_SUBJECTS`` / ``VAL_SUBJECTS`` lists. + +Bring Your Own Data +------------------- +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioMotion4D/`` layout produced by Tutorial 8, not at the repository +``data/`` directory. Edit them to match your own data location. + +Data Required +------------- +Run Tutorial 8 first so +``D:/PhysioMotion4D/duke_data/fitted_kcl_meshes/pm00??/`` contains: + + * ``pm00XX_ssm_surface.vtp`` - reference (template) SSM surface + * ``pm00XX_ssm_pca_coefficients.json`` - fitted PCA coefficient vector + * ``pm00XX_g0TT_ssm_surface.vtp`` - SSM surface at each gated phase TT% + +Outputs (under ``OUTPUT_DIR``) +------------------------------ + * ``physicsnemo_stage_model.pt`` - weights + normalization metadata + * ``physicsnemo_stage_model_epoch_*.pt``- intermittent checkpoints + * ``physicsnemo_stage_model_metadata.json``, ``training_losses.json``, + ``training_validation_rmse.{json,csv}`` - logs + * ``OUTPUT_DIR/test_predictions/statistics_*.csv`` and per-subject predicted + surfaces under ``OUTPUT_DIR/pm00XX/`` + +Extra Install Required +---------------------- +PhysicsNeMo is an optional dependency of PhysioMotion4D. Install it with:: + + pip install "physiomotion4d[physicsnemo]" + +PhysicsNeMo itself requires Python >= 3.11. +""" + +# %% +from __future__ import annotations + +import csv +import json +import logging +import shutil +import sys +from collections import defaultdict +from pathlib import Path +from typing import Any, Optional, cast + +import numpy as np +import pyvista as pv +import torch + + +from physiomotion4d.test_tools import TestTools + +try: + from physicsnemo.models.mlp import FullyConnected +except ImportError as exc: # pragma: no cover - import-time guard + raise ImportError( + "Tutorial 9b requires PhysicsNeMo, which is an optional dependency. " + 'Install with: pip install "physiomotion4d[physicsnemo]" ' + "(requires Python >= 3.11).", + ) from exc + + +# 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. Wrapping consistently across +# tutorials also matches the style of tutorial_01. +if __name__ == "__main__": + # %% + TUTORIALS_DIR = Path(__file__).resolve().parent + FITTED_MESHES_DIR = Path("D:/PhysioMotion4D/duke_data/fitted_kcl_meshes") + PCA_MEAN_VTU = Path("D:/PhysioMotion4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + EPOCHS = 10000 + OUTPUT_DIR = TUTORIALS_DIR / "output" + RMSE_LOG_INTERVAL = ( + 500 # epochs between train/val RMSE reports and checkpoint saves + ) + LOSS_LOG_INTERVAL = 50 # epochs between loss-only console updates + BATCH_SIZE = 262144 # mini-batch size; full dataset lives on GPU (95 GB VRAM) + LEARNING_RATE = 1.0e-3 + LAYER_SIZE = ( + 512 # wider than single-subject model to share capacity across subjects + ) + NUM_LAYERS = 6 + # Explicit subject lists for held-out evaluation splits. All remaining + # subjects are used for training. An error is raised if any listed subject + # is not found in FITTED_MESHES_DIR. Set to None to skip that split entirely. + TEST_SUBJECTS: Optional[list[str]] = ["pm0028"] + VAL_SUBJECTS: Optional[list[str]] = ["pm0027"] + # When True, the x/y/z input coordinates are taken from the PCA mean-shape surface + # (same for every subject). Each subject's PCA coefficients + stage then fully + # describe a query; the displacement target is still relative to that subject's own + # SSM surface (Option B). When False, each subject's own ssm_surface.vtp coordinates + # are used as inputs instead. + USE_MEAN_SHAPE_COORDS = True + LOG_LEVEL = logging.INFO + # Leave as None to train from scratch. Set to the path of a prior run's + # "physicsnemo_stage_model.pt" to resume training from those weights. When + # resuming, a new, numbered output directory is always created so the prior + # run's files are never modified. Normalization stats (coordinate / PCA / + # displacement scales) are inherited from the checkpoint so the weights remain + # valid; only the train/val/test subject split may differ. + RESUME_FROM_WEIGHTS: Optional[Path] = None + + def _next_output_dir(base: Path) -> Path: + """Return the next unused sibling of *base* by appending _1, _2, ...""" + if not base.exists(): + return base + n = 1 + while True: + candidate = base.parent / f"{base.name}_{n}" + if not candidate.exists(): + return candidate + n += 1 + + def _gating_stage_from_filename(mesh_file: Path) -> float: + """Extract the normalised cardiac stage [0, 1] from a ``g0TT`` filename stem. + + For example ``pm0002_g050_ssm_surface.vtp`` -> ``0.50``. + """ + stem = mesh_file.stem # e.g. "pm0002_g050_ssm_surface" + for part in stem.split("_"): + if part.startswith("g") and part[1:].isdigit(): + return int(part[1:]) / 100.0 + raise ValueError(f"Cannot parse gating percentage from filename: {mesh_file}") + + def _uncompiled_state_dict(model: torch.nn.Module) -> dict: + """Return the base model's state dict, unwrapping torch.compile if applied.""" + return cast(dict, getattr(model, "_orig_mod", model).state_dict()) + + def _infer_all_points( + model: "FullyConnected", + norm_coords: np.ndarray, + norm_pca: np.ndarray, + stage: float, + displacement_scale: float, + device: "torch.device", + ) -> np.ndarray: + """Run batched inference over all surface points; return raw displacements (mm).""" + n = len(norm_coords) + pca_tile = np.tile(norm_pca, (n, 1)) + stage_col = np.full((n, 1), stage, dtype=np.float32) + pred_inputs = np.hstack([norm_coords, pca_tile, stage_col]) + chunks: list[np.ndarray] = [] + with torch.no_grad(): + for start in range(0, n, BATCH_SIZE): + stop = min(start + BATCH_SIZE, n) + t = torch.from_numpy(pred_inputs[start:stop].astype(np.float32)).to( + device + ) + chunks.append(model(t).cpu().numpy()) + return np.vstack(chunks) * displacement_scale + + def run_tutorial() -> dict[str, Any]: + """Train a single shared PhysicsNeMo model across all subjects and evaluate. + + Each training sample is a surface point with inputs + ``(x, y, z, pca_c1 ... pca_cN, stage)`` and target the cardiac-motion + displacement from that subject's SSM reference surface to the gated phase + (Option B). Coordinates are taken from the PCA mean shape when + ``USE_MEAN_SHAPE_COORDS`` is True, or from each subject's own SSM surface + otherwise. + + Returns + ------- + dict[str, Any] + Per-subject predicted mesh and evaluation paths, plus shared model paths. + """ + fitted_meshes_dir = FITTED_MESHES_DIR + pca_mean_vtu = PCA_MEAN_VTU + epochs = EPOCHS + rmse_log_interval = RMSE_LOG_INTERVAL + loss_log_interval = LOSS_LOG_INTERVAL + batch_size = BATCH_SIZE + learning_rate = LEARNING_RATE + layer_size = LAYER_SIZE + num_layers = NUM_LAYERS + test_subjects = TEST_SUBJECTS + val_subjects = VAL_SUBJECTS + use_mean_shape_coords = USE_MEAN_SHAPE_COORDS + log_level = LOG_LEVEL + resume_from_weights = RESUME_FROM_WEIGHTS + + logging.basicConfig(level=log_level) + + # In test mode, train for only a couple of epochs so the tutorial test + # exercises the full pipeline without a long GPU run. + if TestTools.running_as_test(): + epochs = min(epochs, 2) + logging.info("Test mode: reducing epochs to %d", epochs) + + # When resuming, always write into a fresh numbered directory so prior + # outputs (weights, logs, CSVs) are never overwritten. + if resume_from_weights is not None: + output_dir = _next_output_dir(OUTPUT_DIR) + logging.info( + f"Resuming from {resume_from_weights}; " + f"new output directory: {output_dir}" + ) + else: + output_dir = OUTPUT_DIR + test_output_dir = output_dir / "test_predictions" + + output_dir.mkdir(parents=True, exist_ok=True) + test_output_dir.mkdir(parents=True, exist_ok=True) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # ------------------------------------------------------------------ # + # 1. Optionally load PCA mean-shape surface (shared coordinate space) # + # ------------------------------------------------------------------ # + mean_shape_coords: Optional[np.ndarray] = None + if use_mean_shape_coords: + logging.info(f"Loading PCA mean shape from {pca_mean_vtu}") + mean_vol = pv.read(str(pca_mean_vtu)) + mean_surf = mean_vol.extract_surface(algorithm="dataset_surface") + mean_shape_coords = np.asarray(mean_surf.points, dtype=np.float32) + logging.info(f"Mean shape surface: {len(mean_shape_coords)} points") + + # ------------------------------------------------------------------ # + # 2. Discover and load all subjects # + # ------------------------------------------------------------------ # + subjects: dict[str, dict] = {} + + for subject_dir in sorted(fitted_meshes_dir.glob("pm????")): + sid = subject_dir.name + ref_file = subject_dir / f"{sid}_ssm_surface.vtp" + pca_file = subject_dir / f"{sid}_ssm_pca_coefficients.json" + mesh_files = sorted(subject_dir.glob(f"{sid}_g0*_ssm_surface.vtp")) + + missing = [p for p in (ref_file, pca_file) if not p.exists()] + if missing or len(mesh_files) < 2: + msg = ( + f"Skipping {sid}: missing files {[str(p) for p in missing]}" + if missing + else f"Skipping {sid}: only {len(mesh_files)} gated phase(s) found" + ) + logging.info(msg) + continue + + ref_mesh = pv.read(str(ref_file)) + subjects[sid] = { + "subject_dir": subject_dir, + "ref_mesh": ref_mesh, + "ref_points": np.asarray(ref_mesh.points, dtype=np.float32), + "pca_coeffs": np.array( + json.loads(pca_file.read_text(encoding="utf-8")), dtype=np.float32 + ), + "mesh_files": mesh_files, + } + + if len(subjects) < 3: + raise RuntimeError( + f"Found only {len(subjects)} valid subject(s); need at least 3 for a " + "train / val / test subject split." + ) + + n_pca = next(iter(subjects.values()))["pca_coeffs"].shape[0] + n_mesh_points = next(iter(subjects.values()))["ref_points"].shape[0] + + if ( + use_mean_shape_coords + and mean_shape_coords is not None + and len(mean_shape_coords) != n_mesh_points + ): + raise RuntimeError( + "Mean-shape surface topology mismatch: expected " + f"{n_mesh_points} points but got {len(mean_shape_coords)}." + ) + + # ------------------------------------------------------------------ # + # 3. Validate explicit subject splits and derive training set # + # ------------------------------------------------------------------ # + all_sids = set(subjects.keys()) + + # Normalise None -> empty list for uniform handling below. + test_list = test_subjects if test_subjects is not None else [] + val_list = val_subjects if val_subjects is not None else [] + + unknown_test = [s for s in test_list if s not in all_sids] + unknown_val = [s for s in val_list if s not in all_sids] + if unknown_test or unknown_val: + parts = [] + if unknown_test: + parts.append(f"TEST_SUBJECTS not found: {unknown_test}") + if unknown_val: + parts.append(f"VAL_SUBJECTS not found: {unknown_val}") + raise ValueError( + "Subject(s) listed in split configuration do not exist in " + f"{fitted_meshes_dir}:\n " + "\n ".join(parts) + ) + + overlap = set(test_list) & set(val_list) + if overlap: + raise ValueError( + f"Subject(s) appear in both TEST_SUBJECTS and VAL_SUBJECTS: {sorted(overlap)}" + ) + + test_sids: set[str] = set(test_list) + val_sids: set[str] = set(val_list) + train_sids: set[str] = all_sids - test_sids - val_sids + + if not train_sids: + raise ValueError( + "No subjects remain for training after applying TEST_SUBJECTS and VAL_SUBJECTS." + ) + + logging.info( + f"Subject split - train: {len(train_sids)}, " + f"val: {len(val_sids)}, test: {len(test_sids)}" + ) + logging.info(f" train subjects: {sorted(train_sids)}") + logging.info(f" val subjects: {sorted(val_sids)}") + logging.info(f" test subjects: {sorted(test_sids)}") + + # ------------------------------------------------------------------ # + # 4. Normalisation statistics # + # When resuming, inherit stats from the prior checkpoint so that the # + # loaded weights remain valid. Otherwise compute from training data. # + # ------------------------------------------------------------------ # + resume_ckpt: Optional[dict] = None + if resume_from_weights is not None: + logging.info(f"Loading prior weights from {resume_from_weights}") + resume_ckpt = torch.load( + str(resume_from_weights), map_location="cpu", weights_only=True + ) + coordinate_mean = np.array(resume_ckpt["coordinate_mean"], dtype=np.float32) + coordinate_scale = np.array( + resume_ckpt["coordinate_scale"], dtype=np.float32 + ) + pca_mean_vec = np.array(resume_ckpt["pca_mean"], dtype=np.float32) + pca_scale_vec = np.array(resume_ckpt["pca_scale"], dtype=np.float32) + displacement_scale = float(resume_ckpt["displacement_scale"]) + logging.info( + "Reusing normalization statistics from prior checkpoint " + f"(displacement_scale={displacement_scale:.4f} mm)." + ) + else: + # Coordinates: use mean shape if available, else pool training ref surfaces. + if use_mean_shape_coords and mean_shape_coords is not None: + coord_ref = mean_shape_coords + else: + coord_ref = np.vstack([subjects[s]["ref_points"] for s in train_sids]) + coordinate_mean = coord_ref.mean(axis=0) + coordinate_scale = coord_ref.std(axis=0) + coordinate_scale = np.where(coordinate_scale == 0.0, 1.0, coordinate_scale) + + # PCA coefficients: per-dimension stats from training subjects only. + train_pca = np.vstack([subjects[s]["pca_coeffs"] for s in train_sids]) + pca_mean_vec = train_pca.mean(axis=0) + pca_scale_vec = train_pca.std(axis=0) + pca_scale_vec = np.where(pca_scale_vec == 0.0, 1.0, pca_scale_vec) + + # ------------------------------------------------------------------ # + # 5. Build combined training and validation datasets # + # ------------------------------------------------------------------ # + logging.info("Building training and validation datasets ...") + training_inputs: list[np.ndarray] = [] + training_targets: list[np.ndarray] = [] + val_inputs: list[np.ndarray] = [] + val_targets: list[np.ndarray] = [] + + def _build_rows( + files: list[Path], + norm_coords: np.ndarray, + pca_tile: np.ndarray, + ref_pts: np.ndarray, + ) -> tuple[list[np.ndarray], list[np.ndarray]]: + inp_rows, tgt_rows = [], [] + for mesh_file in files: + mesh = pv.read(str(mesh_file)) + if mesh.n_points != n_mesh_points: + raise ValueError( + f"{mesh_file} has {mesh.n_points} points, expected {n_mesh_points}." + ) + stage = _gating_stage_from_filename(mesh_file) + stage_col = np.full((len(norm_coords), 1), stage, dtype=np.float32) + inp_rows.append(np.hstack([norm_coords, pca_tile, stage_col])) + tgt_rows.append(np.asarray(mesh.points, dtype=np.float32) - ref_pts) + return inp_rows, tgt_rows + + for sid in sorted(train_sids): + data = subjects[sid] + coords = ( + mean_shape_coords + if use_mean_shape_coords and mean_shape_coords is not None + else data["ref_points"] + ) + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (data["pca_coeffs"] - pca_mean_vec) / pca_scale_vec + pca_tile = np.tile(norm_pca, (n_mesh_points, 1)) + rows_in, rows_tgt = _build_rows( + data["mesh_files"], norm_coords, pca_tile, data["ref_points"] + ) + training_inputs.extend(rows_in) + training_targets.extend(rows_tgt) + + for sid in sorted(val_sids): + data = subjects[sid] + coords = ( + mean_shape_coords + if use_mean_shape_coords and mean_shape_coords is not None + else data["ref_points"] + ) + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (data["pca_coeffs"] - pca_mean_vec) / pca_scale_vec + pca_tile = np.tile(norm_pca, (n_mesh_points, 1)) + rows_in, rows_tgt = _build_rows( + data["mesh_files"], norm_coords, pca_tile, data["ref_points"] + ) + val_inputs.extend(rows_in) + val_targets.extend(rows_tgt) + + inputs_array = np.vstack(training_inputs).astype(np.float32) + targets_array = np.vstack(training_targets).astype(np.float32) + if resume_ckpt is None: + # Fresh run: derive displacement_scale from the training targets. + displacement_scale = float(np.max(np.abs(targets_array))) + if displacement_scale == 0.0: + displacement_scale = 1.0 + targets_array /= displacement_scale + + has_val = len(val_inputs) > 0 + if has_val: + val_inputs_array = np.vstack(val_inputs).astype(np.float32) + val_targets_array = ( + np.vstack(val_targets).astype(np.float32) / displacement_scale + ) + else: + logging.info("No validation subjects configured; skipping val RMSE.") + + logging.info( + f"Training set: {len(inputs_array):,} rows, " + f"val set: {len(val_inputs) if has_val else 0:,} rows, " + f"in_features={inputs_array.shape[1]}, " + f"displacement_scale={displacement_scale:.4f} mm" + ) + + # ------------------------------------------------------------------ # + # 5. Train single shared model # + # ------------------------------------------------------------------ # + in_features = 3 + n_pca + 1 # xyz + pca_coefficients + stage + model = FullyConnected( + in_features=in_features, + layer_size=layer_size, + out_features=3, + num_layers=num_layers, + activation_fn="silu", + skip_connections=True, + ).to(device) + if resume_ckpt is not None: + state = resume_ckpt.get("model_state_dict", resume_ckpt) + model.load_state_dict(state) + logging.info("Loaded model weights from prior checkpoint.") + + # torch.compile requires Triton, which is Linux-only; skip silently on Windows. + if sys.platform != "win32": + try: + model = torch.compile(model) + logging.info( + "torch.compile enabled - first epoch slower while JIT warms up." + ) + except Exception as _compile_err: + logging.info( + f"torch.compile skipped ({_compile_err}); running in eager mode." + ) + else: + logging.info("torch.compile skipped on Windows (Triton unavailable).") + + # Dataset (~3.5 GB) fits comfortably in the 95 GB GPU; keep it resident on GPU + # to eliminate per-batch CPU->GPU transfers (the primary utilization bottleneck). + logging.info("Moving training and validation tensors to GPU ...") + inputs_tensor = torch.from_numpy(inputs_array).to(device) + targets_tensor = torch.from_numpy(targets_array).to(device) + if has_val: + val_inputs_tensor = torch.from_numpy(val_inputs_array).to(device) + val_targets_tensor = torch.from_numpy(val_targets_array).to(device) + n_train = len(inputs_tensor) + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + loss_fn = torch.nn.MSELoss() + + def _batched_rmse_mm(inputs: torch.Tensor, targets: torch.Tensor) -> float: + """Euclidean RMSE in mm, computed in batches (tensors already on GPU).""" + total_sq = 0.0 + n_total = 0 + with torch.no_grad(): + for start in range(0, len(inputs), batch_size): + stop = min(start + batch_size, len(inputs)) + err_mm = ( + model(inputs[start:stop]) - targets[start:stop] + ) * displacement_scale + total_sq += float(torch.sum(err_mm**2)) + n_total += stop - start + return float(np.sqrt(total_sq / n_total)) + + losses: list[float] = [] + rmse_log: list[dict] = [] + + for epoch in range(epochs): + model.train() + epoch_loss = 0.0 + # Shuffle indices on GPU - avoids any CPU involvement mid-epoch. + perm = torch.randperm(n_train, device=device) + for start in range(0, n_train, batch_size): + idx = perm[start : start + batch_size] + batch_in = inputs_tensor[idx] + batch_tgt = targets_tensor[idx] + optimizer.zero_grad(set_to_none=True) + # BF16 autocast: Blackwell tensor cores deliver ~2x FP32 throughput on BF16. + # No GradScaler needed - BF16 exponent range is wide enough to skip scaling. + with torch.amp.autocast(device.type, dtype=torch.bfloat16): + loss = loss_fn(model(batch_in), batch_tgt) + loss.backward() + optimizer.step() + epoch_loss += float(loss.detach()) * len(batch_in) + losses.append(epoch_loss / n_train) + + if (epoch + 1) % loss_log_interval == 0: + logging.info( + " epoch %05d/%d loss=%.6f", epoch + 1, epochs, losses[-1] + ) + + if (epoch + 1) % rmse_log_interval == 0 or epoch + 1 == epochs: + model.eval() + train_rmse = _batched_rmse_mm(inputs_tensor, targets_tensor) + val_rmse = ( + _batched_rmse_mm(val_inputs_tensor, val_targets_tensor) + if has_val + else float("nan") + ) + rmse_log.append( + { + "epoch": epoch + 1, + "train_rmse_mm": train_rmse, + "val_rmse_mm": val_rmse, + } + ) + ckpt_path = ( + output_dir / f"physicsnemo_stage_model_epoch_{epoch + 1:05d}.pt" + ) + torch.save(_uncompiled_state_dict(model), ckpt_path) + logging.info( + "INTERMITTENT TEST epoch %05d/%d " + "train RMSE=%.4f mm val RMSE=%.4f mm checkpoint=%s", + epoch + 1, + epochs, + train_rmse, + val_rmse, + ckpt_path.name, + ) + model.eval() + + # ------------------------------------------------------------------ # + # 6. Save shared model weights + metadata # + # ------------------------------------------------------------------ # + checkpoint_file = output_dir / "physicsnemo_stage_model.pt" + metadata_file = output_dir / "physicsnemo_stage_model_metadata.json" + losses_file = output_dir / "training_losses.json" + + torch.save( + { + "model_state_dict": _uncompiled_state_dict(model), + "in_features": in_features, + "layer_size": layer_size, + "num_layers": num_layers, + "coordinate_mean": coordinate_mean.tolist(), + "coordinate_scale": coordinate_scale.tolist(), + "pca_mean": pca_mean_vec.tolist(), + "pca_scale": pca_scale_vec.tolist(), + "displacement_scale": displacement_scale, + "use_mean_shape_coords": use_mean_shape_coords, + "pca_mean_vtu": str(pca_mean_vtu) if use_mean_shape_coords else None, + "train_subject_ids": sorted(train_sids), + "val_subject_ids": sorted(val_sids), + "test_subject_ids": sorted(test_sids), + "resumed_from": str(resume_from_weights) + if resume_from_weights + else None, + }, + checkpoint_file, + ) + input_feature_names = ( + ( + ["mean_shape_x", "mean_shape_y", "mean_shape_z"] + if use_mean_shape_coords + else ["ssm_x", "ssm_y", "ssm_z"] + ) + + [f"pca_c{i + 1}" for i in range(n_pca)] + + ["stage"] + ) + metadata_file.write_text( + json.dumps( + { + "architecture": "physicsnemo.models.mlp.FullyConnected", + "input_features": input_feature_names, + "output_features": ["dx", "dy", "dz"], + "n_subjects": len(subjects), + "in_features": in_features, + "layer_size": layer_size, + "num_layers": num_layers, + "epochs": epochs, + "n_mesh_points": n_mesh_points, + "learning_rate": learning_rate, + "coordinate_mean": coordinate_mean.tolist(), + "coordinate_scale": coordinate_scale.tolist(), + "pca_mean": pca_mean_vec.tolist(), + "pca_scale": pca_scale_vec.tolist(), + "displacement_scale": displacement_scale, + "use_mean_shape_coords": use_mean_shape_coords, + "displacement_convention": "Option B: relative to subject ssm_surface", + "resumed_from": str(resume_from_weights) + if resume_from_weights + else None, + }, + indent=2, + ), + encoding="utf-8", + ) + losses_file.write_text(json.dumps(losses, indent=2), encoding="utf-8") + rmse_log_file = output_dir / "training_validation_rmse.json" + rmse_log_file.write_text(json.dumps(rmse_log, indent=2), encoding="utf-8") + rmse_csv_file = output_dir / "training_validation_rmse.csv" + with rmse_csv_file.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter( + fh, fieldnames=["epoch", "train_rmse_mm", "val_rmse_mm"] + ) + writer.writeheader() + writer.writerows(rmse_log) + + # Copy PCA assets so the output directory is self-contained for replay. + pca_src_dir = pca_mean_vtu.parent + shutil.copy2(pca_mean_vtu, output_dir / pca_mean_vtu.name) + pca_model_src = pca_src_dir / "pca_model.json" + if pca_model_src.exists(): + shutil.copy2(pca_model_src, output_dir / pca_model_src.name) + # Save the extracted mean-shape surface alongside the volume mesh. + if use_mean_shape_coords and mean_shape_coords is not None: + mean_surf.save(str(output_dir / "pca_mean_surface.vtp")) + + logging.info(f"Shared model saved to {checkpoint_file}") + + # ------------------------------------------------------------------ # + # ------------------------------------------------------------------ # + # 7. Evaluate test and val subjects: all phases -> output/pm00XX/ # + # ------------------------------------------------------------------ # + def _evaluate_subject(sid: str, split_label: str) -> tuple[list[dict], Path]: + data = subjects[sid] + coords = ( + mean_shape_coords + if use_mean_shape_coords and mean_shape_coords is not None + else data["ref_points"] + ) + norm_coords_full = (coords - coordinate_mean) / coordinate_scale + norm_pca = (data["pca_coeffs"] - pca_mean_vec) / pca_scale_vec + + subj_out_dir = output_dir / sid + subj_out_dir.mkdir(parents=True, exist_ok=True) + + sq_err_sum = np.zeros(n_mesh_points, dtype=np.float64) + stats = [] + + for phase_file in data["mesh_files"]: + stage = _gating_stage_from_filename(phase_file) + pred_disps = _infer_all_points( + model, norm_coords_full, norm_pca, stage, displacement_scale, device + ) + pred_points = data["ref_points"] + pred_disps + + pred_mesh = data["ref_mesh"].copy(deep=True) + pred_mesh.points = pred_points + gating_tag = phase_file.stem.split("_ssm_surface")[0].split("_")[-1] + pred_mesh.save( + subj_out_dir / f"{sid}_{gating_tag}_ssm_surface_pred.vtp" + ) + + actual_points = np.asarray( + pv.read(str(phase_file)).points, dtype=np.float32 + ) + errors = pred_points - actual_points + euclidean = np.linalg.norm(errors, axis=1) + sq_err_sum += euclidean.astype(np.float64) ** 2 + stats.append( + { + "subject_id": sid, + "split": split_label, + "gating_tag": gating_tag, + "stage": stage, + "n_points": len(euclidean), + "mean_error_mm": float(euclidean.mean()), + "median_error_mm": float(np.median(euclidean)), + "max_error_mm": float(euclidean.max()), + "rms_error_mm": float(np.sqrt(np.mean(euclidean**2))), + "std_error_mm": float(euclidean.std()), + "mean_abs_error_x_mm": float(np.abs(errors[:, 0]).mean()), + "mean_abs_error_y_mm": float(np.abs(errors[:, 1]).mean()), + "mean_abs_error_z_mm": float(np.abs(errors[:, 2]).mean()), + } + ) + logging.info( + f"{sid} [{split_label}] {gating_tag}: " + f"mean={stats[-1]['mean_error_mm']:.3f} mm " + f"max={stats[-1]['max_error_mm']:.3f} mm" + ) + + point_rmse = np.sqrt(sq_err_sum / len(data["mesh_files"])).astype( + np.float32 + ) + rmse_mesh = data["ref_mesh"].copy(deep=True) + rmse_mesh.point_data["RMSE_mm"] = point_rmse + rmse_file = subj_out_dir / f"{sid}_ssm_surface_rmse.vtp" + rmse_mesh.save(rmse_file) + logging.info( + f"{sid} [{split_label}] per-point RMSE: " + f"mean={point_rmse.mean():.3f} mm max={point_rmse.max():.3f} mm" + ) + return stats, rmse_file + + all_stats = [] + tutorial_outputs = {} + + for sid in sorted(test_sids): + stats, rmse_file = _evaluate_subject(sid, "test") + all_stats.extend(stats) + tutorial_outputs[sid] = { + "split": "test", + "rmse_file": rmse_file, + "final_loss": losses[-1], + "n_phases": len(subjects[sid]["mesh_files"]), + } + + for sid in sorted(val_sids): + stats, rmse_file = _evaluate_subject(sid, "val") + all_stats.extend(stats) + tutorial_outputs[sid] = { + "split": "val", + "rmse_file": rmse_file, + "final_loss": losses[-1], + "n_phases": len(subjects[sid]["mesh_files"]), + } + + # ------------------------------------------------------------------ # + # 8. Write CSV statistics (split column distinguishes test vs val) # + # ------------------------------------------------------------------ # + if all_stats: + per_phase_csv = test_output_dir / "statistics_per_phase.csv" + with per_phase_csv.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(all_stats[0].keys())) + writer.writeheader() + writer.writerows(all_stats) + logging.info(f"Per-phase statistics -> {per_phase_csv}") + + subject_rows = defaultdict(list) + for row in all_stats: + subject_rows[row["subject_id"]].append(row) + + summary_rows = [ + { + "subject_id": sid, + "split": rows[0]["split"], + "n_phases": len(rows), + "mean_error_mm": float(np.mean([r["mean_error_mm"] for r in rows])), + "mean_max_error_mm": float( + np.mean([r["max_error_mm"] for r in rows]) + ), + "overall_max_error_mm": float( + np.max([r["max_error_mm"] for r in rows]) + ), + "mean_rms_error_mm": float( + np.mean([r["rms_error_mm"] for r in rows]) + ), + } + for sid, rows in sorted(subject_rows.items()) + ] + summary_csv = test_output_dir / "statistics_summary.csv" + with summary_csv.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(summary_rows[0].keys())) + writer.writeheader() + writer.writerows(summary_rows) + logging.info(f"Summary statistics -> {summary_csv}") + + return tutorial_outputs + + # %% + # Run this cell in VS Code or Cursor: + tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py b/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py new file mode 100644 index 0000000..b0ba8d5 --- /dev/null +++ b/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py @@ -0,0 +1,534 @@ +""" +Tutorial 10a (MGN): Predict cardiac stage meshes for a subject using a trained +PhysicsNeMo MeshGraphNet. + +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +Loads the MeshGraphNet checkpoint trained by Tutorial 9a +(``tutorial_09a_cardiac_train_physicsnemo_mgn.py``) and predicts cardiac +surfaces for one subject. Can be run from the command line, or cell-by-cell / +as a tutorial test via the ``run_tutorial`` entry point (which uses the +``DEFAULT_SUBJECT`` / ``DEFAULT_EPOCH`` constants below). + +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioMotion4D/`` layout and the Tutorial 9a run directory, not at the +repository ``data/`` directory. + +Usage (command line) +-------------------- + py tutorial_10a_cardiac_eval_physicsnemo_mgn.py pm0028 --epoch 1500 --out results/pm0028_mgn + py tutorial_10a_cardiac_eval_physicsnemo_mgn.py pm0028 --epoch 1500 --out results/pm0028_mgn --stages 0.0 0.25 0.5 0.75 + +Arguments + subject Subject ID, e.g. pm0028 + --epoch Training epoch whose checkpoint to load, e.g. 1500 + --out Output directory (created if missing) + --stages RR-interval fractions to predict (0-1). If omitted, predicts at + every gated phase found in the subject's fitted-meshes directory + and computes error statistics against the ground-truth surfaces. + When supplied, only the requested stages are predicted (no error + stats, since ground-truth surfaces may not exist). + +Reads weights from OUTPUT_DIR/mgn_stage_model_epoch_EEEEE.pt and normalization +metadata from OUTPUT_DIR/mgn_stage_model.pt, both produced by Tutorial 9a. The +shared mesh graph is replayed from OUTPUT_DIR/shared_edge_index.pt and +OUTPUT_DIR/shared_edge_features.pt. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import logging +import sys +from pathlib import Path +from typing import Any, Optional, cast + +import numpy as np +import pyvista as pv +import torch + +try: + from torch_geometric.data import Data + + from physicsnemo.models.meshgraphnet import MeshGraphNet +except ImportError as exc: + raise ImportError( + "Requires PhysicsNeMo and PyTorch Geometric. Install with:\n" + ' pip install "physiomotion4d[physicsnemo]"\n' + " pip install torch-geometric" + ) from exc + +logger = logging.getLogger("tutorial_10a_cardiac_eval_physicsnemo_mgn") + +TUTORIALS_DIR = Path(__file__).resolve().parent +FITTED_MESHES_DIR = Path("D:/PhysioMotion4D/duke_data/fitted_kcl_meshes") +# Tutorial 9a run directory to evaluate (matches that trainer's OUTPUT_DIR). +OUTPUT_DIR = TUTORIALS_DIR / "output_mgn" + +# Defaults used by run_tutorial() when this file is run with no CLI arguments +# (e.g. cell-by-cell execution or as a tutorial test). +DEFAULT_SUBJECT = "pm0028" +DEFAULT_EPOCH = 1500 +DEFAULT_OUT_DIR = TUTORIALS_DIR / "output_mgn" / "eval_mgn" / DEFAULT_SUBJECT + +# These match tutorial_09a_cardiac_train_physicsnemo_mgn.py. Older checkpoints +# store only processor_size and hidden_dim in metadata, so keep the layer +# defaults explicit. +DEFAULT_NUM_LAYERS_PROCESSOR = 2 +DEFAULT_NUM_LAYERS_ENCODER = 2 +DEFAULT_NUM_LAYERS_DECODER = 2 +DEFAULT_NUM_PROCESSOR_CHECKPOINT_SEGMENTS = 0 + + +def _latest_epoch_checkpoint(output_dir: Path) -> Optional[int]: + """Return the highest epoch number among saved checkpoints, or None if none exist.""" + epochs = [] + for ckpt in output_dir.glob("mgn_stage_model_epoch_*.pt"): + try: + epochs.append(int(ckpt.stem.rsplit("_", 1)[-1])) + except ValueError: + continue + return max(epochs) if epochs else None + + +def _gating_stage_from_filename(mesh_file: Path) -> float: + for part in mesh_file.stem.split("_"): + if part.startswith("g") and part[1:].isdigit(): + return int(part[1:]) / 100.0 + raise ValueError(f"Cannot parse gating percentage from filename: {mesh_file}") + + +def _state_dict_from_checkpoint(checkpoint: Any) -> dict[str, torch.Tensor]: + state = ( + checkpoint.get("model_state_dict", checkpoint) + if isinstance(checkpoint, dict) + else checkpoint + ) + if not isinstance(state, dict): + raise TypeError( + f"Checkpoint did not contain a state dict: {type(checkpoint)!r}" + ) + return state + + +def _strip_compile_prefix(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + if all(k.startswith("_orig_mod.") for k in state): + return {k.removeprefix("_orig_mod."): v for k, v in state.items()} + return state + + +def _build_model(meta: dict[str, Any], device: torch.device) -> MeshGraphNet: + return MeshGraphNet( + input_dim_nodes=int(meta["in_features"]), + input_dim_edges=4, + output_dim=3, + processor_size=int(meta["processor_size"]), + hidden_dim_processor=int(meta["hidden_dim"]), + hidden_dim_node_encoder=int(meta["hidden_dim"]), + num_layers_node_encoder=int( + meta.get("num_layers_node_encoder", DEFAULT_NUM_LAYERS_ENCODER) + ), + hidden_dim_node_decoder=int(meta["hidden_dim"]), + num_layers_node_decoder=int( + meta.get("num_layers_node_decoder", DEFAULT_NUM_LAYERS_DECODER) + ), + hidden_dim_edge_encoder=int(meta["hidden_dim"]), + num_layers_edge_processor=int( + meta.get("num_layers_edge_processor", DEFAULT_NUM_LAYERS_PROCESSOR) + ), + num_layers_node_processor=int( + meta.get("num_layers_node_processor", DEFAULT_NUM_LAYERS_PROCESSOR) + ), + aggregation="mean", + num_processor_checkpoint_segments=int( + meta.get( + "num_processor_checkpoint_segments", + DEFAULT_NUM_PROCESSOR_CHECKPOINT_SEGMENTS, + ) + ), + ).to(device) + + +def _infer_all_points( + model: MeshGraphNet, + norm_coords: np.ndarray, + norm_pca: np.ndarray, + stage: float, + shared_graph: Data, + shared_edge_feats: torch.Tensor, + displacement_scale: float, + device: torch.device, +) -> np.ndarray: + n = len(norm_coords) + pca_tile = np.tile(norm_pca, (n, 1)) + stage_col = np.full((n, 1), stage, dtype=np.float32) + node_feats = torch.tensor( + np.hstack([norm_coords, pca_tile, stage_col]), + dtype=torch.float32, + device=device, + ) + with torch.no_grad(): + pred = model(node_feats, shared_edge_feats, shared_graph) + return np.asarray(pred.cpu().numpy()) * displacement_scale + + +def predict( + subject: str, + epoch: int, + out_dir: Path, + stages: Optional[list[float]] = None, +) -> dict[str, Any]: + subject_dir = FITTED_MESHES_DIR / subject + ref_file = subject_dir / f"{subject}_ssm_surface.vtp" + pca_file = subject_dir / f"{subject}_ssm_pca_coefficients.json" + + for p in (ref_file, pca_file): + if not p.exists(): + sys.exit(f"Missing: {p}") + + if stages is None: + phase_files = sorted(subject_dir.glob(f"{subject}_g0*_ssm_surface.vtp")) + if not phase_files: + sys.exit(f"No gated phase files found in {subject_dir}") + else: + bad = [s for s in stages if not 0.0 <= s <= 1.0] + if bad: + sys.exit(f"--stages values must be in [0, 1]; got: {bad}") + phase_files = [] + + meta_ckpt = OUTPUT_DIR / "mgn_stage_model.pt" + epoch_ckpt = OUTPUT_DIR / f"mgn_stage_model_epoch_{epoch:05d}.pt" + edge_index_file = OUTPUT_DIR / "shared_edge_index.pt" + edge_features_file = OUTPUT_DIR / "shared_edge_features.pt" + for p in (meta_ckpt, epoch_ckpt, edge_index_file, edge_features_file): + if not p.exists(): + sys.exit(f"Missing trained GNN artifact: {p}") + + meta = torch.load(meta_ckpt, map_location="cpu", weights_only=True) + coordinate_mean = np.array(meta["coordinate_mean"], dtype=np.float32) + coordinate_scale = np.array(meta["coordinate_scale"], dtype=np.float32) + pca_mean_vec = np.array(meta["pca_mean"], dtype=np.float32) + pca_scale_vec = np.array(meta["pca_scale"], dtype=np.float32) + displacement_scale = float(meta["displacement_scale"]) + use_mean_shape_coords = bool(meta["use_mean_shape_coords"]) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _build_model(meta, device) + epoch_state = torch.load(epoch_ckpt, map_location=device, weights_only=True) + model.load_state_dict( + _strip_compile_prefix(_state_dict_from_checkpoint(epoch_state)) + ) + model.eval() + + shared_edge_index = torch.load( + edge_index_file, map_location=device, weights_only=True + ) + shared_edge_feats = torch.load( + edge_features_file, map_location=device, weights_only=True + ) + n_graph_nodes = int(shared_edge_index.max().item()) + 1 + shared_graph = Data(edge_index=shared_edge_index, num_nodes=n_graph_nodes).to( + device + ) + shared_edge_feats = shared_edge_feats.to(device) + + ref_mesh = cast(pv.PolyData, pv.read(str(ref_file))) + ref_points = np.asarray(ref_mesh.points, dtype=np.float32) + pca_coeffs = np.array( + json.loads(pca_file.read_text(encoding="utf-8")), dtype=np.float32 + ) + + if use_mean_shape_coords: + mean_surf_vtp = OUTPUT_DIR / "pca_mean_surface.vtp" + if not mean_surf_vtp.exists(): + sys.exit(f"Mean-shape surface not found: {mean_surf_vtp}") + coords = np.asarray(pv.read(str(mean_surf_vtp)).points, dtype=np.float32) + else: + coords = ref_points + + if len(coords) != len(ref_points): + sys.exit( + f"Topology mismatch: coordinate surface has {len(coords)} points, " + f"but {ref_file} has {len(ref_points)} points." + ) + if len(coords) != n_graph_nodes: + sys.exit( + f"Graph topology mismatch: graph has {n_graph_nodes} nodes, " + f"but coordinate surface has {len(coords)} points." + ) + if pca_coeffs.shape != pca_mean_vec.shape: + sys.exit( + f"PCA coefficient mismatch: subject has shape {pca_coeffs.shape}, " + f"checkpoint expects {pca_mean_vec.shape}." + ) + + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (pca_coeffs - pca_mean_vec) / pca_scale_vec + + out_dir.mkdir(parents=True, exist_ok=True) + logger.info("Subject: %s epoch: %s device: %s", subject, epoch, device) + + if stages is None: + return _predict_with_errors( + subject, + epoch, + ref_mesh, + ref_points, + norm_coords, + norm_pca, + phase_files, + model, + shared_graph, + shared_edge_feats, + displacement_scale, + device, + out_dir, + ) + return _predict_arbitrary_stages( + subject, + ref_mesh, + ref_points, + norm_coords, + norm_pca, + stages, + model, + shared_graph, + shared_edge_feats, + displacement_scale, + device, + out_dir, + ) + + +def _predict_with_errors( + subject: str, + epoch: int, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + norm_coords: np.ndarray, + norm_pca: np.ndarray, + phase_files: list[Path], + model: MeshGraphNet, + shared_graph: Data, + shared_edge_feats: torch.Tensor, + displacement_scale: float, + device: torch.device, + out_dir: Path, +) -> dict[str, Any]: + n_points = len(ref_points) + sq_err_sum = np.zeros(n_points, dtype=np.float64) + results: list[dict[str, Any]] = [] + predicted_files: list[Path] = [] + + for phase_file in phase_files: + stage = _gating_stage_from_filename(phase_file) + pred_disps = _infer_all_points( + model, + norm_coords, + norm_pca, + stage, + shared_graph, + shared_edge_feats, + displacement_scale, + device, + ) + pred_points = (ref_points + pred_disps).astype(np.float32) + actual_points = np.asarray(pv.read(str(phase_file)).points, dtype=np.float32) + error_vec = pred_points - actual_points + error_mag = np.linalg.norm(error_vec, axis=1).astype(np.float32) + sq_err_sum += error_mag.astype(np.float64) ** 2 + gating_tag = phase_file.stem.split("_ssm_surface")[0].split("_")[-1] + results.append( + { + "gating_tag": gating_tag, + "stage": stage, + "pred_points": pred_points, + "error_vec": error_vec, + "error_mag": error_mag, + } + ) + + point_rmse = np.sqrt(sq_err_sum / len(results)).astype(np.float32) + + stats_rows: list[dict[str, Any]] = [] + for r in results: + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = r["pred_points"] + pred_mesh.point_data["error_x"] = r["error_vec"][:, 0] + pred_mesh.point_data["error_y"] = r["error_vec"][:, 1] + pred_mesh.point_data["error_z"] = r["error_vec"][:, 2] + pred_mesh.point_data["error_mm"] = r["error_mag"] + pred_mesh.point_data["rmse_mm"] = point_rmse + + tag = r["gating_tag"] + out_file = out_dir / f"{subject}_{tag}_ssm_surface_pred.vtp" + pred_mesh.save(str(out_file)) + predicted_files.append(out_file) + logger.info(" %s", out_file.name) + + em = r["error_mag"] + stats_rows.append( + { + "subject": subject, + "epoch": epoch, + "gating_tag": tag, + "stage": r["stage"], + "n_points": n_points, + "mean_error_mm": float(em.mean()), + "median_error_mm": float(np.median(em)), + "rms_error_mm": float(np.sqrt(np.mean(em**2))), + "std_error_mm": float(em.std()), + "max_error_mm": float(em.max()), + "mean_abs_error_x_mm": float(np.abs(r["error_vec"][:, 0]).mean()), + "mean_abs_error_y_mm": float(np.abs(r["error_vec"][:, 1]).mean()), + "mean_abs_error_z_mm": float(np.abs(r["error_vec"][:, 2]).mean()), + } + ) + + stats_rows.append( + { + "subject": subject, + "epoch": epoch, + "gating_tag": "ALL", + "stage": "", + "n_points": n_points, + "mean_error_mm": float(np.mean([r["mean_error_mm"] for r in stats_rows])), + "median_error_mm": float( + np.mean([r["median_error_mm"] for r in stats_rows]) + ), + "rms_error_mm": float( + np.sqrt(sq_err_sum.sum() / (n_points * len(results))) + ), + "std_error_mm": float(np.mean([r["std_error_mm"] for r in stats_rows])), + "max_error_mm": float(np.max([r["max_error_mm"] for r in stats_rows])), + "mean_abs_error_x_mm": float( + np.mean([r["mean_abs_error_x_mm"] for r in stats_rows]) + ), + "mean_abs_error_y_mm": float( + np.mean([r["mean_abs_error_y_mm"] for r in stats_rows]) + ), + "mean_abs_error_z_mm": float( + np.mean([r["mean_abs_error_z_mm"] for r in stats_rows]) + ), + } + ) + + csv_path = out_dir / "statistics.csv" + with csv_path.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(stats_rows[0].keys())) + writer.writeheader() + writer.writerows(stats_rows) + + summary = next(r for r in stats_rows if r["gating_tag"] == "ALL") + logger.info( + "%d predictions -> %s (overall RMS=%.4f mm, mean=%.4f mm, max=%.4f mm, " + "statistics -> %s)", + len(results), + out_dir, + summary["rms_error_mm"], + summary["mean_error_mm"], + summary["max_error_mm"], + csv_path.name, + ) + return { + "subject": subject, + "epoch": epoch, + "predicted_files": predicted_files, + "statistics_csv": csv_path, + "summary": summary, + } + + +def _predict_arbitrary_stages( + subject: str, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + norm_coords: np.ndarray, + norm_pca: np.ndarray, + stages: list[float], + model: MeshGraphNet, + shared_graph: Data, + shared_edge_feats: torch.Tensor, + displacement_scale: float, + device: torch.device, + out_dir: Path, +) -> dict[str, Any]: + predicted_files: list[Path] = [] + for stage in stages: + pred_disps = _infer_all_points( + model, + norm_coords, + norm_pca, + stage, + shared_graph, + shared_edge_feats, + displacement_scale, + device, + ) + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = (ref_points + pred_disps).astype(np.float32) + tag = f"s{round(stage * 100):03d}" + out_file = out_dir / f"{subject}_{tag}_ssm_surface_pred.vtp" + pred_mesh.save(str(out_file)) + predicted_files.append(out_file) + logger.info(" %s", out_file.name) + + logger.info("%d predictions -> %s", len(stages), out_dir) + return { + "subject": subject, + "predicted_files": predicted_files, + "statistics_csv": None, + "summary": None, + } + + +def run_tutorial() -> dict[str, Any]: + """Tutorial / test entry point: evaluate DEFAULT_SUBJECT at the latest checkpoint. + + Used when the script is run with no command-line arguments (cell-by-cell + execution or as a tutorial test). Picks the highest-numbered checkpoint + under OUTPUT_DIR so this works whether Tutorial 9a ran a full training + pass or the reduced test-mode epoch count; falls back to DEFAULT_EPOCH if + no checkpoints are found. Returns the prediction outputs dict. + """ + epoch = _latest_epoch_checkpoint(OUTPUT_DIR) or DEFAULT_EPOCH + return predict(DEFAULT_SUBJECT, epoch, DEFAULT_OUT_DIR) + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Predict cardiac stage meshes for one subject with MeshGraphNet.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Example: py tutorial_10a_cardiac_eval_physicsnemo_mgn.py " + "pm0028 --epoch 1500 --out results/pm0028_mgn" + ), + ) + ap.add_argument("subject", help="Subject ID, e.g. pm0028") + ap.add_argument( + "--epoch", type=int, required=True, help="Training epoch, e.g. 1500" + ) + ap.add_argument("--out", type=Path, required=True, help="Output directory") + ap.add_argument( + "--stages", + type=float, + nargs="+", + metavar="FRAC", + help=( + "RR-interval fractions to predict, e.g. --stages 0.0 0.25 0.5 0.75 " + "(omit to predict at every existing gated phase with error statistics)" + ), + ) + args = ap.parse_args() + predict(args.subject, args.epoch, args.out, stages=args.stages) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + if len(sys.argv) > 1: + main() + else: + # %% + # Tutorial / test entry point (no CLI arguments) + tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py b/tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py new file mode 100644 index 0000000..d3c9bb6 --- /dev/null +++ b/tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py @@ -0,0 +1,469 @@ +""" +Tutorial 10b (MLP): Predict cardiac stage meshes for a subject using a trained +PhysicsNeMo MLP checkpoint. + +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +Loads the MLP checkpoint trained by Tutorial 9b +(``tutorial_09b_cardiac_train_physicsnemo_mlp.py``) and predicts cardiac +surfaces for one subject. Can be run from the command line, or cell-by-cell / +as a tutorial test via the ``run_tutorial`` entry point (which uses the +``DEFAULT_SUBJECT`` / ``DEFAULT_EPOCH`` constants below). + +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioMotion4D/`` layout and the Tutorial 9b run directory, not at the +repository ``data/`` directory. + +Usage (command line) +-------------------- + py tutorial_10b_cardiac_eval_physicsnemo_mlp.py pm0002 --epoch 5000 --out results/pm0002 + py tutorial_10b_cardiac_eval_physicsnemo_mlp.py pm0002 --epoch 5000 --out results/pm0002 --stages 0.0 0.25 0.5 0.75 + +Arguments + subject Subject ID, e.g. pm0002 + --epoch Training epoch whose checkpoint to load, e.g. 5000 + --out Output directory (created if missing) + --stages RR-interval fractions to predict (0-1). If omitted, predicts at + every gated phase found in the subject's fitted-meshes directory + and computes error statistics against the ground-truth surfaces. + When supplied, only the requested stages are predicted (no error + stats, since ground-truth surfaces may not exist). + +Reads weights from OUTPUT_DIR/physicsnemo_stage_model_epoch_EEEEE.pt +and normalisation metadata from OUTPUT_DIR/physicsnemo_stage_model.pt, +both produced by Tutorial 9b. (The ``OUTPUT_DIR`` constant below selects which +run directory to evaluate.) + +Outputs (--stages omitted) +-------------------------- +One predicted .vtp surface per gated phase, each carrying per-point arrays: + + error_x/y/z signed offset (pred - target) in mm + error_mm Euclidean distance from target in mm + rmse_mm per-point RMSE over all stages (same value in every file) + +statistics.csv per-stage and summary error statistics + +Outputs (--stages supplied) +--------------------------- +One predicted .vtp surface per requested stage; no error arrays or CSV. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import logging +import sys +from pathlib import Path +from typing import Any, Optional, cast + +import numpy as np +import pyvista as pv +import torch + +try: + from physicsnemo.models.mlp import FullyConnected +except ImportError as exc: + raise ImportError( + 'Requires PhysicsNeMo. Install with: pip install "physiomotion4d[physicsnemo]"' + ) from exc + +logger = logging.getLogger("tutorial_10b_cardiac_eval_physicsnemo_mlp") + +TUTORIALS_DIR = Path(__file__).resolve().parent +FITTED_MESHES_DIR = Path("D:/PhysioMotion4D/duke_data/fitted_kcl_meshes") +# Tutorial 9b run directory to evaluate (matches that trainer's OUTPUT_DIR). +OUTPUT_DIR = TUTORIALS_DIR / "output" +BATCH_SIZE = 262144 + +# Defaults used by run_tutorial() when this file is run with no CLI arguments +# (e.g. cell-by-cell execution or as a tutorial test). +DEFAULT_SUBJECT = "pm0028" +DEFAULT_EPOCH = 5000 +DEFAULT_OUT_DIR = TUTORIALS_DIR / "output" / "eval_mlp" / DEFAULT_SUBJECT + + +def _latest_epoch_checkpoint(output_dir: Path) -> Optional[int]: + """Return the highest epoch number among saved checkpoints, or None if none exist.""" + epochs = [] + for ckpt in output_dir.glob("physicsnemo_stage_model_epoch_*.pt"): + try: + epochs.append(int(ckpt.stem.rsplit("_", 1)[-1])) + except ValueError: + continue + return max(epochs) if epochs else None + + +def _gating_stage_from_filename(mesh_file: Path) -> float: + for part in mesh_file.stem.split("_"): + if part.startswith("g") and part[1:].isdigit(): + return int(part[1:]) / 100.0 + raise ValueError(f"Cannot parse gating percentage from filename: {mesh_file}") + + +def _infer( + model: "FullyConnected", + norm_coords: np.ndarray, + pca_tile: np.ndarray, + stage: float, + displacement_scale: float, + device: "torch.device", +) -> np.ndarray: + n = len(norm_coords) + stage_col = np.full((n, 1), stage, dtype=np.float32) + pred_inputs = np.hstack([norm_coords, pca_tile, stage_col]) + chunks: list[np.ndarray] = [] + with torch.no_grad(): + for start in range(0, n, BATCH_SIZE): + stop = min(start + BATCH_SIZE, n) + t = torch.from_numpy(pred_inputs[start:stop]).to(device) + chunks.append(model(t).cpu().numpy()) + return np.vstack(chunks) * displacement_scale + + +def predict( + subject: str, + epoch: int, + out_dir: Path, + stages: Optional[list[float]] = None, +) -> dict[str, Any]: + subject_dir = FITTED_MESHES_DIR / subject + ref_file = subject_dir / f"{subject}_ssm_surface.vtp" + pca_file = subject_dir / f"{subject}_ssm_pca_coefficients.json" + + for p in (ref_file, pca_file): + if not p.exists(): + sys.exit(f"Missing: {p}") + + # When --stages is not given, discover gated phase files and validate them. + if stages is None: + phase_files = sorted(subject_dir.glob(f"{subject}_g0*_ssm_surface.vtp")) + if not phase_files: + sys.exit(f"No gated phase files found in {subject_dir}") + else: + bad = [s for s in stages if not 0.0 <= s <= 1.0] + if bad: + sys.exit(f"--stages values must be in [0, 1]; got: {bad}") + phase_files = [] + + meta_ckpt = OUTPUT_DIR / "physicsnemo_stage_model.pt" + if not meta_ckpt.exists(): + sys.exit( + f"Metadata checkpoint not found: {meta_ckpt}\n" + "Run Tutorial 9b (tutorial_09b_cardiac_train_physicsnemo_mlp.py) first." + ) + meta = torch.load(meta_ckpt, map_location="cpu", weights_only=True) + + epoch_ckpt = OUTPUT_DIR / f"physicsnemo_stage_model_epoch_{epoch:05d}.pt" + if not epoch_ckpt.exists(): + sys.exit(f"Epoch checkpoint not found: {epoch_ckpt}") + + coordinate_mean = np.array(meta["coordinate_mean"], dtype=np.float32) + coordinate_scale = np.array(meta["coordinate_scale"], dtype=np.float32) + pca_mean_vec = np.array(meta["pca_mean"], dtype=np.float32) + pca_scale_vec = np.array(meta["pca_scale"], dtype=np.float32) + displacement_scale = float(meta["displacement_scale"]) + use_mean_shape_coords: bool = meta["use_mean_shape_coords"] + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = FullyConnected( + in_features=int(meta["in_features"]), + layer_size=int(meta["layer_size"]), + out_features=3, + num_layers=int(meta["num_layers"]), + activation_fn="silu", + skip_connections=True, + ).to(device) + model.load_state_dict( + torch.load(epoch_ckpt, map_location=device, weights_only=True) + ) + model.eval() + + ref_mesh = cast(pv.PolyData, pv.read(str(ref_file))) + ref_points = np.asarray(ref_mesh.points, dtype=np.float32) + pca_coeffs = np.array( + json.loads(pca_file.read_text(encoding="utf-8")), dtype=np.float32 + ) + + if use_mean_shape_coords: + mean_surf_vtp = OUTPUT_DIR / "pca_mean_surface.vtp" + pca_mean_vtu_path = meta.get("pca_mean_vtu") + if mean_surf_vtp.exists(): + coords = np.asarray(pv.read(str(mean_surf_vtp)).points, dtype=np.float32) + elif pca_mean_vtu_path and Path(pca_mean_vtu_path).exists(): + vol = pv.read(str(pca_mean_vtu_path)) + coords = np.asarray( + vol.extract_surface(algorithm="dataset_surface").points, + dtype=np.float32, + ) + else: + sys.exit( + f"Mean-shape surface not found. Expected {mean_surf_vtp} " + f"or {pca_mean_vtu_path}" + ) + else: + coords = ref_points + + if len(coords) != len(ref_points): + sys.exit( + f"Topology mismatch: coordinate surface has {len(coords)} points, " + f"but {ref_file} has {len(ref_points)} points." + ) + if pca_coeffs.shape != pca_mean_vec.shape: + sys.exit( + f"PCA coefficient mismatch: subject has shape {pca_coeffs.shape}, " + f"checkpoint expects {pca_mean_vec.shape}." + ) + + norm_coords = (coords - coordinate_mean) / coordinate_scale + norm_pca = (pca_coeffs - pca_mean_vec) / pca_scale_vec + pca_tile = np.tile(norm_pca, (len(norm_coords), 1)) + + out_dir.mkdir(parents=True, exist_ok=True) + logger.info("Subject: %s epoch: %s device: %s", subject, epoch, device) + + n_points = len(ref_points) + + if stages is None: + return _predict_with_errors( + subject, + epoch, + ref_mesh, + ref_points, + norm_coords, + pca_tile, + displacement_scale, + phase_files, + model, + device, + out_dir, + n_points, + ) + return _predict_arbitrary_stages( + subject, + ref_mesh, + ref_points, + norm_coords, + pca_tile, + displacement_scale, + stages, + model, + device, + out_dir, + ) + + +def _predict_with_errors( + subject: str, + epoch: int, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + norm_coords: np.ndarray, + pca_tile: np.ndarray, + displacement_scale: float, + phase_files: list[Path], + model: "FullyConnected", + device: "torch.device", + out_dir: Path, + n_points: int, +) -> dict[str, Any]: + """Predict at each existing gated phase; compute and embed error statistics.""" + sq_err_sum = np.zeros(n_points, dtype=np.float64) + results: list[dict] = [] + predicted_files: list[Path] = [] + + for phase_file in phase_files: + stage = _gating_stage_from_filename(phase_file) + pred_disps = _infer( + model, norm_coords, pca_tile, stage, displacement_scale, device + ) + pred_points = (ref_points + pred_disps).astype(np.float32) + actual_points = np.asarray(pv.read(str(phase_file)).points, dtype=np.float32) + error_vec = pred_points - actual_points + error_mag = np.linalg.norm(error_vec, axis=1).astype(np.float32) + sq_err_sum += error_mag.astype(np.float64) ** 2 + gating_tag = phase_file.stem.split("_ssm_surface")[0].split("_")[-1] + results.append( + { + "gating_tag": gating_tag, + "stage": stage, + "pred_points": pred_points, + "error_vec": error_vec, + "error_mag": error_mag, + } + ) + + point_rmse = np.sqrt(sq_err_sum / len(results)).astype(np.float32) + + stats_rows: list[dict] = [] + for r in results: + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = r["pred_points"] + pred_mesh.point_data["error_x"] = r["error_vec"][:, 0] + pred_mesh.point_data["error_y"] = r["error_vec"][:, 1] + pred_mesh.point_data["error_z"] = r["error_vec"][:, 2] + pred_mesh.point_data["error_mm"] = r["error_mag"] + pred_mesh.point_data["rmse_mm"] = point_rmse + + tag = r["gating_tag"] + out_file = out_dir / f"{subject}_{tag}_ssm_surface_pred.vtp" + pred_mesh.save(str(out_file)) + predicted_files.append(out_file) + logger.info(" %s", out_file.name) + + em = r["error_mag"] + stats_rows.append( + { + "subject": subject, + "epoch": epoch, + "gating_tag": tag, + "stage": r["stage"], + "n_points": n_points, + "mean_error_mm": float(em.mean()), + "median_error_mm": float(np.median(em)), + "rms_error_mm": float(np.sqrt(np.mean(em**2))), + "std_error_mm": float(em.std()), + "max_error_mm": float(em.max()), + "mean_abs_error_x_mm": float(np.abs(r["error_vec"][:, 0]).mean()), + "mean_abs_error_y_mm": float(np.abs(r["error_vec"][:, 1]).mean()), + "mean_abs_error_z_mm": float(np.abs(r["error_vec"][:, 2]).mean()), + } + ) + + stats_rows.append( + { + "subject": subject, + "epoch": epoch, + "gating_tag": "ALL", + "stage": "", + "n_points": n_points, + "mean_error_mm": float(np.mean([r["mean_error_mm"] for r in stats_rows])), + "median_error_mm": float( + np.mean([r["median_error_mm"] for r in stats_rows]) + ), + "rms_error_mm": float( + np.sqrt(sq_err_sum.sum() / (n_points * len(results))) + ), + "std_error_mm": float(np.mean([r["std_error_mm"] for r in stats_rows])), + "max_error_mm": float(np.max([r["max_error_mm"] for r in stats_rows])), + "mean_abs_error_x_mm": float( + np.mean([r["mean_abs_error_x_mm"] for r in stats_rows]) + ), + "mean_abs_error_y_mm": float( + np.mean([r["mean_abs_error_y_mm"] for r in stats_rows]) + ), + "mean_abs_error_z_mm": float( + np.mean([r["mean_abs_error_z_mm"] for r in stats_rows]) + ), + } + ) + + csv_path = out_dir / "statistics.csv" + with csv_path.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(stats_rows[0].keys())) + writer.writeheader() + writer.writerows(stats_rows) + + summary = next(r for r in stats_rows if r["gating_tag"] == "ALL") + logger.info( + "%d predictions -> %s (overall RMS=%.4f mm, mean=%.4f mm, max=%.4f mm, " + "statistics -> %s)", + len(results), + out_dir, + summary["rms_error_mm"], + summary["mean_error_mm"], + summary["max_error_mm"], + csv_path.name, + ) + return { + "subject": subject, + "epoch": epoch, + "predicted_files": predicted_files, + "statistics_csv": csv_path, + "summary": summary, + } + + +def _predict_arbitrary_stages( + subject: str, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + norm_coords: np.ndarray, + pca_tile: np.ndarray, + displacement_scale: float, + stages: list[float], + model: "FullyConnected", + device: "torch.device", + out_dir: Path, +) -> dict[str, Any]: + """Predict at caller-specified RR-interval fractions; no ground-truth comparison.""" + predicted_files: list[Path] = [] + for stage in stages: + pred_disps = _infer( + model, norm_coords, pca_tile, stage, displacement_scale, device + ) + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = (ref_points + pred_disps).astype(np.float32) + tag = f"s{round(stage * 100):03d}" + out_file = out_dir / f"{subject}_{tag}_ssm_surface_pred.vtp" + pred_mesh.save(str(out_file)) + predicted_files.append(out_file) + logger.info(" %s", out_file.name) + + logger.info("%d predictions -> %s", len(stages), out_dir) + return { + "subject": subject, + "predicted_files": predicted_files, + "statistics_csv": None, + "summary": None, + } + + +def run_tutorial() -> dict[str, Any]: + """Tutorial / test entry point: evaluate DEFAULT_SUBJECT at the latest checkpoint. + + Used when the script is run with no command-line arguments (cell-by-cell + execution or as a tutorial test). Picks the highest-numbered checkpoint + under OUTPUT_DIR so this works whether Tutorial 9b ran a full training + pass or the reduced test-mode epoch count; falls back to DEFAULT_EPOCH if + no checkpoints are found. Returns the prediction outputs dict. + """ + epoch = _latest_epoch_checkpoint(OUTPUT_DIR) or DEFAULT_EPOCH + return predict(DEFAULT_SUBJECT, epoch, DEFAULT_OUT_DIR) + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Predict cardiac stage meshes for one subject.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Example: py tutorial_10b_cardiac_eval_physicsnemo_mlp.py " + "pm0002 --epoch 5000 --out results/pm0002" + ), + ) + ap.add_argument("subject", help="Subject ID, e.g. pm0002") + ap.add_argument( + "--epoch", type=int, required=True, help="Training epoch, e.g. 5000" + ) + ap.add_argument("--out", type=Path, required=True, help="Output directory") + ap.add_argument( + "--stages", + type=float, + nargs="+", + metavar="FRAC", + help="RR-interval fractions to predict, e.g. --stages 0.0 0.25 0.5 0.75 " + "(omit to predict at every existing gated phase with error statistics)", + ) + args = ap.parse_args() + predict(args.subject, args.epoch, args.out, stages=args.stages) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + if len(sys.argv) > 1: + main() + else: + # %% + # Tutorial / test entry point (no CLI arguments) + tutorial_results = run_tutorial()