diff --git a/docs/architecture.rst b/docs/architecture.rst index 92065dc..951a644 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -91,34 +91,36 @@ Primary Workflows AI Surrogate Workflows (PhysicsNeMo) ===================================== -The final tier of tutorials (``tutorials/tutorial_08`` through -``tutorial_10``) turns a fitted statistical shape model into a trained AI +The final tier of tutorials (``tutorials/tutorial_08cd`` through +``tutorial_10d``) turns a fitted statistical shape model into a trained AI physiological surrogate, replacing the explicit per-phase registration solve with a learned model at inference time: -``tutorial_08_cardiac_fit_model.py`` +``tutorial_08cd_byod_fit_model_to_patients.py`` Fits the cardiac PCA model to a patient (via ``WorkflowFitStatisticalModelToPatient``) and propagates the fitted mesh through every gated phase using ICON-based registration (``WorkflowReconstructHighres4DCT``), producing the per-phase mesh/surface pairs used as AI surrogate training data. -``tutorial_09a_cardiac_train_physicsnemo_mgn.py`` / -``tutorial_09b_cardiac_train_physicsnemo_mlp.py`` - Train a PhysicsNeMo surrogate — a graph-based ``MeshGraphNet`` - (``physicsnemo.models.meshgraphnet``) or a fully connected MLP — on the - Tutorial 8 output to predict per-phase cardiac mesh deformation directly - from the fitted SSM coefficients. Requires the ``[physicsnemo]`` extra - (and ``torch-geometric`` for the MeshGraphNet variant); Python >= 3.11. - -``tutorial_10a_cardiac_eval_physicsnemo_mgn.py`` / -``tutorial_10b_cardiac_eval_physicsnemo_mlp.py`` - Load a trained MeshGraphNet or MLP checkpoint and predict/score cardiac - surfaces without running registration, i.e. the AI surrogate stands in for - ``WorkflowReconstructHighres4DCT`` at inference time. - -These tutorials are not wrapped in a ``physiotwin4d`` workflow class today — -they call the PhysicsNeMo model classes directly — but they follow the same +``tutorial_09c_byod_train_physicsnemo_mgn.py`` / +``tutorial_09d_byod_train_physicsnemo_mlp.py`` + Train a PhysicsNeMo surrogate — a graph-based ``MeshGraphNet`` via + ``WorkflowTrainPhysicsNeMoMGN`` or a fully connected MLP via + ``WorkflowTrainPhysicsNeMoMLP`` — on the Tutorial 8cd output to predict + per-phase cardiac mesh deformation directly from the fitted SSM + coefficients. Requires the ``[physicsnemo]`` extra (and ``torch-geometric`` + for the MeshGraphNet variant); Python >= 3.11. + +``tutorial_10c_byod_eval_physicsnemo_mgn.py`` / +``tutorial_10d_byod_eval_physicsnemo_mlp.py`` + Load a trained MeshGraphNet or MLP checkpoint (via + ``WorkflowInferPhysicsNeMoMGN`` / ``WorkflowInferPhysicsNeMoMLP``) and + predict/score cardiac surfaces without running registration, i.e. the AI + surrogate stands in for ``WorkflowReconstructHighres4DCT`` at inference time. + +These tutorials are thin drivers over the ``WorkflowTrainPhysicsNeMo`` / +``WorkflowInferPhysicsNeMo`` workflow classes; they follow the same fit -> propagate -> train -> predict pattern the rest of the workflow layer uses, and are the intended template for future cardiac, respiratory, and electrophysiology AI surrogates. diff --git a/docs/index.rst b/docs/index.rst index 42fc318..caa1e97 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -22,8 +22,8 @@
- - 01 + + 1a

Heart-Gated CT to Animated USD

Convert cardiac 4D CT frames into registered contours and an animated OpenUSD model.

Slicer-Heart-CT @@ -34,23 +34,23 @@

Segment one CT phase and export patient anatomy as VTK PolyData surfaces.

Slicer-Heart-CT - + 03 +

VTK Surface Series to Animated USD

+

Convert VTK meshes into a time-sampled USD scene for Omniverse playback.

+ Tutorial 2 output + + + 4a

Create a PCA Shape Model

Build a statistical shape model from aligned cardiac meshes.

KCL-Heart-Model - - 04 + + 5a

Fit Statistical Model to Patient

Fit a PCA heart model to patient-specific anatomy for model-based reconstruction.

- Tutorial 3 output - - - 05 -

VTK Surface Series to Animated USD

-

Convert VTK meshes into a time-sampled USD scene for Omniverse playback.

- Tutorial 2 output + Tutorial 4a output 06 @@ -58,23 +58,23 @@

Register respiratory CT phases and reconstruct a higher-resolution 4D volume series.

DirLab-4DCT - - 08 + + 8cd

Fit the Cardiac SSM and Propagate Through Gated Phases

Fit a PCA heart model to the reference phase and propagate it to every gated phase with ICON registration.

Bring your own cardiac data - - 09 + + 9cd

Train a PhysicsNeMo Cardiac Stage Model

-

Train a PhysicsNeMo MeshGraphNet (9a) or MLP (9b) to predict cardiac meshes at requested stages.

- Tutorial 8 output +

Train a PhysicsNeMo MeshGraphNet (9c) or MLP (9d) to predict cardiac meshes at requested stages.

+ Tutorial 8cd output - - 10 + + 10cd

Predict and Evaluate Cardiac Surfaces

-

Load a Tutorial 9 checkpoint and predict cardiac surfaces at gated phases or caller-specified stages.

- Tutorial 9a / 9b output +

Load a Tutorial 9c/9d checkpoint and predict cardiac surfaces at gated phases or caller-specified stages.

+ Tutorial 9c / 9d output
diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 6413afb..17e2d8b 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -20,7 +20,7 @@ The ``tutorials/`` directory contains eleven end-to-end scripts covering nine major workflows (Tutorials 9 and 10 each have MeshGraphNet and MLP variants). Each script is a ``# %%`` percent-cell Python script that exercises the workflow classes directly. Run as a regular file -(``python tutorials/tutorial_01_...py``) or cell-by-cell in VS Code or Cursor. +(``python tutorials/tutorial_01a_...py``) or cell-by-cell in VS Code or Cursor. See :doc:`tutorials` for the NVIDIA-styled tutorial card index, dataset requirements, script paths, and workflow details. @@ -41,7 +41,7 @@ tutorials: .. code-block:: bash - python tutorials/tutorial_01_heart_gated_ct_to_usd.py + python tutorials/tutorial_01a_heart_gated_ct_to_usd.py python tutorials/tutorial_02_ct_to_vtk.py @@ -52,16 +52,16 @@ recommended run order. Recommended run order: -1. Tutorials 1 and 2 first, after downloading Slicer-Heart-CT data. -2. Tutorial 5 after Tutorial 2 (consumes Tutorial 2 output). -3. Tutorial 3 after downloading KCL-Heart-Model. -4. Tutorial 4 after Tutorial 3 because it can consume Tutorial 3 output. +1. Tutorials 1a and 2 first, after downloading Slicer-Heart-CT data. +2. Tutorial 3 after Tutorial 2 (consumes Tutorial 2 output). +3. Tutorial 4a after downloading KCL-Heart-Model. +4. Tutorial 5a after Tutorial 4a because it can consume Tutorial 4a output. 5. Tutorial 6 after downloading DirLab-4DCT (manual). -6. Tutorial 8 after preparing your own cardiac gated CT, labelmaps, KCL volume +6. Tutorial 8cd after preparing your own cardiac gated CT, labelmaps, KCL volume PCA model, and ICON weights (bring-your-own-data). -7. Tutorial 9a and/or 9b after Tutorial 8 because they train from its fitted +7. Tutorial 9c and/or 9d after Tutorial 8cd because they train from its fitted meshes. -8. Tutorial 10a and/or 10b after Tutorial 9a / 9b because they evaluate the +8. Tutorial 10c and/or 10d after Tutorial 9c / 9d because they evaluate the trained checkpoints. Prerequisites diff --git a/docs/tutorials.rst b/docs/tutorials.rst index a617f2e..61e17eb 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -19,8 +19,8 @@ Tutorials
- - 01 + + 1a

Heart-Gated CT to Animated USD

Convert cardiac 4D CT frames into registered contours and an animated OpenUSD model.

Slicer-Heart-CT @@ -31,23 +31,23 @@ Tutorials

Segment one CT phase and export patient anatomy as VTK PolyData surfaces.

Slicer-Heart-CT - + 03 +

VTK Surface Series to Animated USD

+

Convert VTK meshes into a time-sampled USD scene for Omniverse playback.

+ Tutorial 2 output + + + 4a

Create a PCA Shape Model

Build a statistical shape model from aligned cardiac meshes.

KCL-Heart-Model - - 04 + + 5a

Fit Statistical Model to Patient

Fit a PCA heart model to patient-specific anatomy for model-based reconstruction.

- Tutorial 3 output - - - 05 -

VTK Surface Series to Animated USD

-

Convert VTK meshes into a time-sampled USD scene for Omniverse playback.

- Tutorial 2 output + Tutorial 4a output 06 @@ -55,23 +55,23 @@ Tutorials

Register respiratory CT phases and reconstruct a higher-resolution 4D volume series.

DirLab-4DCT - - 08 + + 8cd

Fit the Cardiac SSM and Propagate Through Gated Phases

Fit a PCA heart model to the reference phase and propagate it to every gated phase with ICON registration.

Bring your own cardiac data - - 09 + + 9cd

Train a PhysicsNeMo Cardiac Stage Model

-

Train a PhysicsNeMo MeshGraphNet (9a) or MLP (9b) to predict cardiac meshes at requested stages.

- Tutorial 8 output +

Train a PhysicsNeMo MeshGraphNet (9c) or MLP (9d) to predict cardiac meshes at requested stages.

+ Tutorial 8cd output - - 10 + + 10cd

Predict and Evaluate Cardiac Surfaces

-

Load a Tutorial 9 checkpoint and predict cardiac surfaces at gated phases or caller-specified stages.

- Tutorial 9a / 9b output +

Load a Tutorial 9c/9d checkpoint and predict cardiac surfaces at gated phases or caller-specified stages.

+ Tutorial 9c / 9d output
@@ -86,15 +86,15 @@ Before running any tutorial below, fetch its sample dataset with the physiotwin4d-download-data Slicer-Heart-CT --directory data/Slicer-Heart-CT physiotwin4d-download-data KCL-Heart-Model --directory data/KCL-Heart-Model -This covers the data used by Tutorials 1-5. Run +This covers the data used by Tutorials 1a-5a. Run ``physiotwin4d-download-data --help`` for all options, and see :doc:`cli_scripts/download_data` for dataset sizes, source URLs, and directory layouts. ``DirLab-4DCT`` (used by Tutorial 6) is **not** auto-downloaded — DIR-Lab distributes it manually and may require registration; see -``data/DirLab-4DCT/README.md`` for instructions. Tutorials 8-10 are -bring-your-own-data (see the note before Tutorial 8) and do not use a +``data/DirLab-4DCT/README.md`` for instructions. Tutorials 8cd-10cd are +bring-your-own-data (see the note before Tutorial 8cd) and do not use a downloadable sample. Recommended Run Order @@ -106,23 +106,23 @@ directories by default. Edit those constants for tutorial exploration, or use the installed ``physiotwin4d-*`` CLI commands when you need command-line path arguments. -1. Run Tutorials 1 and 2 after downloading Slicer-Heart-CT data (above). -2. Run Tutorial 5 after Tutorial 2 because it consumes Tutorial 2 output. -3. Run Tutorial 3 after downloading KCL-Heart-Model (above). -4. Run Tutorial 4 after Tutorial 3 because it can consume the PCA model output. +1. Run Tutorials 1a and 2 after downloading Slicer-Heart-CT data (above). +2. Run Tutorial 3 after Tutorial 2 because it consumes Tutorial 2 output. +3. Run Tutorial 4a after downloading KCL-Heart-Model (above). +4. Run Tutorial 5a after Tutorial 4a because it can consume the PCA model output. 5. Run Tutorial 6 after downloading DirLab-4DCT (manual — see above). -6. Run Tutorial 8 after preparing your own cardiac gated CT, labelmaps, KCL +6. Run Tutorial 8cd 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 +7. Run Tutorial 9c and/or 9d after Tutorial 8cd because they train from its fitted meshes. -8. Run Tutorial 10a and/or 10b after Tutorial 9a / 9b because they evaluate +8. Run Tutorial 10c and/or 10d after Tutorial 9c / 9d because they evaluate the trained checkpoints. -Tutorial 1: Heart-Gated CT to Animated USD -========================================== +Tutorial 1a: Heart-Gated CT to Animated USD +=========================================== Script - ``tutorials/tutorial_01_heart_gated_ct_to_usd.py`` + ``tutorials/tutorial_01a_heart_gated_ct_to_usd.py`` Workflow ``WorkflowConvertImageToUSD`` @@ -168,7 +168,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_01_heart_gated_ct_to_usd.py + python tutorials/tutorial_01a_heart_gated_ct_to_usd.py Outputs Registered phase images, transformed contours, preview screenshots, and an @@ -232,11 +232,11 @@ Run Outputs Segmentation artifacts, VTK PolyData surfaces, and preview screenshots. -Tutorial 3: Create a PCA Shape Model -==================================== +Tutorial 4a: Create a PCA Shape Model +===================================== Script - ``tutorials/tutorial_03_create_statistical_model.py`` + ``tutorials/tutorial_04a_heart_create_statistical_model.py`` Workflow ``WorkflowCreateStatisticalModel`` @@ -273,16 +273,16 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_03_create_statistical_model.py + python tutorials/tutorial_04a_heart_create_statistical_model.py Outputs PCA model files, mean shape, and component diagnostics. -Tutorial 4: Fit Statistical Model to Patient -============================================ +Tutorial 5a: Fit Statistical Model to Patient +============================================= Script - ``tutorials/tutorial_04_fit_statistical_model_to_patient.py`` + ``tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py`` Workflow ``WorkflowFitStatisticalModelToPatient`` @@ -321,16 +321,16 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_04_fit_statistical_model_to_patient.py + python tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py Outputs Patient-fitted statistical model surfaces and registration diagnostics. -Tutorial 5: VTK Surface Series to Animated USD +Tutorial 3: VTK Surface Series to Animated USD ============================================== Script - ``tutorials/tutorial_05_vtk_to_usd.py`` + ``tutorials/tutorial_03_vtk_to_usd.py`` Workflow ``WorkflowConvertVTKToUSD`` @@ -378,7 +378,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_05_vtk_to_usd.py + python tutorials/tutorial_03_vtk_to_usd.py Outputs Time-sampled USD scene and conversion logs for Omniverse inspection. @@ -439,19 +439,18 @@ Outputs .. 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:/PhysioTwin4D/`` 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. + Tutorials 8cd-10cd form the cardiac mesh stage-prediction pipeline and are + **bring-your-own-data**: unlike the earlier data-driven tutorials they do not + use the repository ``data/`` directory or a downloadable sample. Their path + constants point at a local ``D:/PhysioTwin4D/`` cardiac layout (gated CT, + labelmaps, the KCL volume PCA model, and ICON weights); edit those constants + to match your own data. -Tutorial 8: Fit the Cardiac SSM and Propagate Through Gated Phases -================================================================== +Tutorial 8cd: Fit the Cardiac SSM and Propagate Through Gated Phases +=================================================================== Script - ``tutorials/tutorial_08_cardiac_fit_model.py`` + ``tutorials/tutorial_08cd_byod_fit_model_to_patients.py`` Workflow ``WorkflowFitStatisticalModelToPatient`` (PCA registration) and @@ -508,30 +507,32 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_08_cardiac_fit_model.py + python tutorials/tutorial_08cd_byod_fit_model_to_patients.py Outputs Per-patient fitted SSM mesh/surface, PCA coefficients, and the SSM warped to every gated phase, all written under ``OUTPUT_DIR``. -Tutorial 9a / 9b: Train a PhysicsNeMo Cardiac Stage Model +Tutorial 9c / 9d: Train a PhysicsNeMo Cardiac Stage Model ========================================================= Script - ``tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py`` (MeshGraphNet) and - ``tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py`` (MLP) + ``tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py`` (MLP) Inner API usage - Unlike Tutorials 1-8, these do not build a ``physiotwin4d`` workflow — - they import a PhysicsNeMo model class directly and train it on Tutorial 8's - fitted meshes: ``physicsnemo.models.meshgraphnet.MeshGraphNet`` (9a) or - ``physicsnemo.models.mlp.FullyConnected`` (9b). They are the intended - template for future cardiac, respiratory, and electrophysiology AI - surrogates, following the same fit -> propagate -> train -> predict - pattern as the rest of the workflow layer. + These are thin drivers over the reusable + ``WorkflowTrainPhysicsNeMoMGN`` (9c) and ``WorkflowTrainPhysicsNeMoMLP`` (9d) + workflows, which train ``physicsnemo.models.meshgraphnet.MeshGraphNet`` or + ``physicsnemo.models.mlp.FullyConnected`` on Tutorial 8cd's fitted meshes. The + tutorials discover subjects, write one JSON manifest per subject, split them + into train / validation / held-out test, and call ``process()``. They are the + intended template for future cardiac, respiratory, and electrophysiology AI + surrogates, following the same fit -> propagate -> train -> predict pattern as + the rest of the workflow layer. Dataset - Tutorial 8 fitted-mesh outputs. + Tutorial 8cd fitted-mesh outputs. Preview .. figure:: assets/example.gif @@ -554,28 +555,29 @@ Extra install Run .. code-block:: bash - python tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py - python tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py + python tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py + python tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py Outputs Shared PhysicsNeMo checkpoints, training metadata, loss / RMSE histories, and held-out predictions written under each trainer's ``OUTPUT_DIR``. -Tutorial 10a / 10b: Predict and Evaluate Cardiac Surfaces +Tutorial 10c / 10d: Predict and Evaluate Cardiac Surfaces ========================================================= Script - ``tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py`` (MeshGraphNet) and - ``tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py`` (MLP) + ``tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py`` (MLP) Inner API usage - Loads a Tutorial 9 checkpoint and predicts cardiac surfaces for one - subject at each gated phase (with error statistics) or at - caller-specified stages — the AI surrogate standing in for - ``WorkflowReconstructHighres4DCT`` at inference time. + Thin drivers over ``WorkflowInferPhysicsNeMoMGN`` / ``WorkflowInferPhysicsNeMoMLP`` + that load a Tutorial 9c/9d checkpoint and predict cardiac surfaces for one + subject at each gated phase (with error statistics) or at caller-specified + stages — the AI surrogate standing in for ``WorkflowReconstructHighres4DCT`` + at inference time. Dataset - Tutorial 9a / 9b trained checkpoints plus the Tutorial 8 fitted meshes. + Tutorial 9c / 9d trained checkpoints plus the Tutorial 8cd fitted meshes. Preview .. figure:: assets/example.gif @@ -593,14 +595,14 @@ Preview Run .. code-block:: bash - python tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py pm0002 --epoch 5000 --out results/pm0002 + python tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py pm0002 --out results/pm0002 Run with no arguments to use the ``run_tutorial`` entry point and its - ``DEFAULT_SUBJECT`` / ``DEFAULT_EPOCH`` constants. + ``DEFAULT_SUBJECT`` / ``DEFAULT_OUT_DIR`` constants. Outputs - Predicted ``.vtp`` surfaces per phase (with per-point error arrays when - ground truth exists) and a ``statistics.csv`` error summary. + Predicted ``.vtp`` surfaces per phase, a per-point ``RMSE_mm`` surface, and a + ``statistics_per_phase.csv`` error summary (when ground-truth phases exist). Dataset Notes ============= diff --git a/pyproject.toml b/pyproject.toml index ec945e1..9df685b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,11 +110,14 @@ cuda13 = [ "torchaudio", ] # PhysicsNeMo pulls in a large dependency stack (CUDA-only toolchain, narrower -# Python support). It is only needed for Tutorial 9 (mesh stage MLP). Install -# with ``pip install physiotwin4d[physicsnemo]`` when you need it. Note: +# Python support). It is needed for the PhysicsNeMo train/infer workflows and +# Tutorial 9 (mesh stage MLP / MeshGraphNet). ``torch-geometric`` supplies the +# graph batching / connectivity utilities used by the MeshGraphNet (MGN) path. +# Install with ``pip install physiotwin4d[physicsnemo]`` when you need it. Note: # nvidia-physicsnemo itself requires Python >= 3.11. physicsnemo = [ "nvidia-physicsnemo>=2.0.0", + "torch-geometric>=2.5.0", ] dev = [ "bumpver>=2023.0.0", @@ -173,6 +176,8 @@ physiotwin4d-convert-vtk-to-usd = "physiotwin4d.cli.convert_vtk_to_usd:main" physiotwin4d-create-statistical-model = "physiotwin4d.cli.create_statistical_model:main" physiotwin4d-download-data = "physiotwin4d.cli.download_data:main" physiotwin4d-fit-statistical-model-to-patient = "physiotwin4d.cli.fit_statistical_model_to_patient:main" +physiotwin4d-train-physicsnemo = "physiotwin4d.cli.train_physicsnemo:main" +physiotwin4d-infer-physicsnemo = "physiotwin4d.cli.infer_physicsnemo:main" physiotwin4d-reconstruct-highres-4d-ct = "physiotwin4d.cli.reconstruct_highres_4d_ct:main" physiotwin4d-visualize-pca-modes = "physiotwin4d.cli.visualize_pca_modes:main" @@ -249,6 +254,8 @@ module = [ "sklearn.*", "torch", "torch.*", + "torch_geometric", + "torch_geometric.*", "totalsegmentator.*", "trimesh", "trimesh.*", @@ -267,11 +274,11 @@ module = [ "physiotwin4d.cli.visualize_pca_modes", "physiotwin4d.vtk_to_usd.mesh_utils", "physiotwin4d.vtk_to_usd.vtk_reader", - "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", + "tutorial_08cd_byod_fit_model_to_patients", + "tutorial_09c_byod_train_physicsnemo_mgn", + "tutorial_09d_byod_train_physicsnemo_mlp", + "tutorial_10c_byod_eval_physicsnemo_mgn", + "tutorial_10d_byod_eval_physicsnemo_mlp", ] disable_error_code = ["import-not-found", "import-untyped"] diff --git a/src/physiotwin4d/__init__.py b/src/physiotwin4d/__init__.py index 9008d4c..e85597c 100644 --- a/src/physiotwin4d/__init__.py +++ b/src/physiotwin4d/__init__.py @@ -89,6 +89,16 @@ from .workflow_fit_statistical_model_to_patient import ( WorkflowFitStatisticalModelToPatient, ) +from .workflow_infer_physicsnemo import ( + WorkflowInferPhysicsNeMo, + WorkflowInferPhysicsNeMoMGN, + WorkflowInferPhysicsNeMoMLP, +) +from .workflow_train_physicsnemo import ( + WorkflowTrainPhysicsNeMo, + WorkflowTrainPhysicsNeMoMGN, + WorkflowTrainPhysicsNeMoMLP, +) __all__ = [ # Workflow classes @@ -99,6 +109,12 @@ "WorkflowFineTuneICONRegistration", "WorkflowReconstructHighres4DCT", "WorkflowFitStatisticalModelToPatient", + "WorkflowTrainPhysicsNeMo", + "WorkflowTrainPhysicsNeMoMGN", + "WorkflowTrainPhysicsNeMoMLP", + "WorkflowInferPhysicsNeMo", + "WorkflowInferPhysicsNeMoMGN", + "WorkflowInferPhysicsNeMoMLP", # Segmentation classes "SegmentAnatomyBase", "SegmentChestTotalSegmentator", diff --git a/src/physiotwin4d/cli/infer_physicsnemo.py b/src/physiotwin4d/cli/infer_physicsnemo.py new file mode 100644 index 0000000..bc58fa8 --- /dev/null +++ b/src/physiotwin4d/cli/infer_physicsnemo.py @@ -0,0 +1,126 @@ +"""Command-line interface for PhysicsNeMo cardiac mesh-stage inference. + +Loads a trained model directory and predicts either from a per-subject manifest +(``--manifest``) or from a PCA shape-parameter file (``--shape-parameters``). The +network is auto-detected from the checkpoint files unless ``--network`` is given. +With ``--reference-image`` a deformation field and surface-normal image are +rasterized onto that image's grid. +""" + +import argparse +import sys +from pathlib import Path + + +def _detect_network(model_dir: Path) -> str: + """Return 'mgn' or 'mlp' based on the checkpoint present in ``model_dir``.""" + for tag in ("mgn", "mlp"): + if (model_dir / f"{tag}_stage_model.pt").exists(): + return tag + raise FileNotFoundError( + f"No _stage_model.pt found in {model_dir}; pass --network explicitly." + ) + + +def main() -> int: + """CLI entry point for PhysicsNeMo inference.""" + parser = argparse.ArgumentParser( + description="Infer cardiac mesh stages with a trained PhysicsNeMo model.", + ) + parser.add_argument("--model-dir", required=True, help="Trained model directory.") + parser.add_argument( + "--network", + choices=("mgn", "mlp", "auto"), + default="auto", + help="Network architecture (auto-detected from the checkpoint by default).", + ) + parser.add_argument("--epoch", type=int, default=None, help="Checkpoint epoch.") + parser.add_argument("--output", default=None, help="Output directory.") + + # Manifest-driven mode. + parser.add_argument("--manifest", default=None, help="Per-subject manifest JSON.") + parser.add_argument( + "--stages", + nargs="*", + type=float, + default=None, + help="Arbitrary stages to predict (manifest mode; omit for phase eval).", + ) + + # Manifest-free single-subject mode. + parser.add_argument( + "--shape-parameters", default=None, help="PCA shape-parameter JSON file." + ) + parser.add_argument( + "--stage", type=float, default=None, help="Target stage (single-subject mode)." + ) + parser.add_argument( + "--ground-truth", default=None, help="Optional ground-truth surface .vtp." + ) + parser.add_argument( + "--reference-image", + default=None, + help="Reference image; when given, write a deformation field + normal image.", + ) + + args = parser.parse_args() + model_dir = Path(args.model_dir) + network = args.network if args.network != "auto" else _detect_network(model_dir) + output = Path(args.output) if args.output else None + + from ..workflow_infer_physicsnemo import ( + WorkflowInferPhysicsNeMo, + WorkflowInferPhysicsNeMoMGN, + WorkflowInferPhysicsNeMoMLP, + ) + + workflow: WorkflowInferPhysicsNeMo + if network == "mgn": + workflow = WorkflowInferPhysicsNeMoMGN( + model_directory=model_dir, epoch=args.epoch + ) + else: + workflow = WorkflowInferPhysicsNeMoMLP( + model_directory=model_dir, epoch=args.epoch + ) + + if args.manifest is not None: + result = workflow.predict( + Path(args.manifest), stages=args.stages, output_directory=output + ) + print(f"Predicted {len(result['predicted_surfaces'])} surface(s).") + return 0 + + if args.shape_parameters is not None: + if args.stage is None: + parser.error("--stage is required with --shape-parameters.") + if args.reference_image is not None: + import itk + + reference_image = itk.imread(args.reference_image) + result = workflow.create_deformation_field( + Path(args.shape_parameters), + args.stage, + reference_image, + output_directory=output, + ) + print( + f"Deformation field written to {result.get('deformation_field_file')}" + ) + return 0 + + ground_truth = Path(args.ground_truth) if args.ground_truth else None + result = workflow.predict_single( + Path(args.shape_parameters), + args.stage, + ground_truth=ground_truth, + output_directory=output, + ) + print(f"Predicted surface written to {result['predicted_surface']}") + return 0 + + parser.error("Provide either --manifest or --shape-parameters.") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/physiotwin4d/cli/train_physicsnemo.py b/src/physiotwin4d/cli/train_physicsnemo.py new file mode 100644 index 0000000..f19acc2 --- /dev/null +++ b/src/physiotwin4d/cli/train_physicsnemo.py @@ -0,0 +1,135 @@ +"""Command-line interface for training a PhysicsNeMo cardiac mesh-stage model. + +Selects the MeshGraphNet (``--network mgn``) or fully connected +(``--network mlp``) trainer, wires the per-subject manifest lists and tuning +options, and runs the workflow. +""" + +import argparse +import sys +from pathlib import Path + + +def _apply_common(workflow: object, args: argparse.Namespace) -> None: + """Apply the shared tuning setters when supplied on the command line.""" + if args.epochs is not None: + workflow.set_epochs(args.epochs) # type: ignore[attr-defined] + if args.batch_size is not None: + workflow.set_batch_size(args.batch_size) # type: ignore[attr-defined] + if args.learning_rate is not None: + workflow.set_learning_rate(args.learning_rate) # type: ignore[attr-defined] + if args.cache_size is not None: + workflow.set_cache_size(args.cache_size) # type: ignore[attr-defined] + + +def main() -> int: + """CLI entry point for PhysicsNeMo training.""" + parser = argparse.ArgumentParser( + description="Train a PhysicsNeMo cardiac mesh-stage model (MGN or MLP).", + ) + parser.add_argument( + "--network", + choices=("mgn", "mlp"), + required=True, + help="Network architecture: mgn (MeshGraphNet) or mlp (FullyConnected).", + ) + parser.add_argument( + "--train-manifest", + nargs="+", + required=True, + metavar="JSON", + help="Per-subject training manifest files.", + ) + parser.add_argument( + "--val-manifest", + nargs="*", + default=[], + metavar="JSON", + help="Per-subject validation manifest files (intermittent RMSE).", + ) + parser.add_argument( + "--pca-mean-mesh", + required=True, + help="PCA template mesh (e.g. pca_mean.vtu) matching pca_model.json.", + ) + parser.add_argument( + "--output", + required=True, + help="Output directory for checkpoints, metadata and logs.", + ) + parser.add_argument( + "--resume-from", + default=None, + help="Optional prior _stage_model.pt to resume from.", + ) + + # Shared tuning. + parser.add_argument("--epochs", type=int, default=None) + parser.add_argument( + "--batch-size", type=int, default=None, help="Mini-batch size in samples." + ) + parser.add_argument("--learning-rate", type=float, default=None) + parser.add_argument( + "--cache-size", + type=int, + default=None, + help="RAM cache budget (decoded phase arrays); 0 = unbounded.", + ) + + # MGN-specific. + parser.add_argument("--processor-size", type=int, default=None) + parser.add_argument("--hidden-dim", type=int, default=None) + # MLP-specific. + parser.add_argument("--layer-size", type=int, default=None) + # Shared architecture depth (both networks expose set_num_layers). + parser.add_argument("--num-layers", type=int, default=None) + + args = parser.parse_args() + + train_manifests = [Path(p) for p in args.train_manifest] + val_manifests = [Path(p) for p in args.val_manifest] + pca_mean_mesh = Path(args.pca_mean_mesh) + output_directory = Path(args.output) + resume_from = Path(args.resume_from) if args.resume_from else None + + if args.network == "mgn": + from ..workflow_train_physicsnemo import WorkflowTrainPhysicsNeMoMGN + + mgn = WorkflowTrainPhysicsNeMoMGN( + train_manifests=train_manifests, + val_manifests=val_manifests, + pca_mean_mesh=pca_mean_mesh, + output_directory=output_directory, + resume_from=resume_from, + ) + _apply_common(mgn, args) + if args.processor_size is not None: + mgn.set_processor_size(args.processor_size) + if args.hidden_dim is not None: + mgn.set_hidden_dim(args.hidden_dim) + if args.num_layers is not None: + mgn.set_num_layers(args.num_layers) + result = mgn.process() + else: + from ..workflow_train_physicsnemo import WorkflowTrainPhysicsNeMoMLP + + mlp = WorkflowTrainPhysicsNeMoMLP( + train_manifests=train_manifests, + val_manifests=val_manifests, + pca_mean_mesh=pca_mean_mesh, + output_directory=output_directory, + resume_from=resume_from, + ) + _apply_common(mlp, args) + if args.layer_size is not None: + mlp.set_layer_size(args.layer_size) + if args.num_layers is not None: + mlp.set_num_layers(args.num_layers) + result = mlp.process() + + print(f"Training complete. Checkpoint: {result['checkpoint']}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/physiotwin4d/physicsnemo_tools.py b/src/physiotwin4d/physicsnemo_tools.py new file mode 100644 index 0000000..1184d61 --- /dev/null +++ b/src/physiotwin4d/physicsnemo_tools.py @@ -0,0 +1,351 @@ +"""Shared helpers for the PhysicsNeMo train / infer workflows. + +This module holds the pieces common to the MeshGraphNet (MGN) and fully +connected (MLP) PhysicsNeMo workflows so the workflow classes stay focused on +orchestration. It provides: + +- :class:`SubjectManifest` / :func:`parse_manifest` — the per-subject JSON + manifest that lists a reference surface, a PCA shape-parameter file, and the + gated-phase surfaces with their stages. +- :func:`build_node_features` — the shared per-vertex feature layout + ``[mean_coords_norm, pca_norm (tiled), stage]`` used by both networks. +- :func:`mesh_to_edge_index` / :func:`compute_edge_features` — MGN mesh-graph + construction from the shared mean-shape surface. +- :func:`reconstruct_reference_points` — rebuild a subject reference surface + from PCA shape parameters (``P = mean + Σ b_i·std_i·eigenvector_i``), used for + manifest-free single-subject inference. +- :func:`uncompiled_state_dict` / :func:`strip_compile_prefix` — checkpoint I/O + that is robust to ``torch.compile`` wrapping. +- :class:`PhaseSampleDataset` — a lazy ``(subject, phase)`` sample provider with + a bounded in-RAM cache so the training set need not fit in memory. + +``torch`` and ``torch_geometric`` are optional dependencies; every function that +needs them imports them locally so ``import physiotwin4d`` works without the +``[physicsnemo]`` extra installed. +""" + +from __future__ import annotations + +import json +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +import pyvista as pv + +if TYPE_CHECKING: # imported lazily at runtime; typed here for mypy only + import torch + + +# --------------------------------------------------------------------------- # +# Per-subject manifest # +# --------------------------------------------------------------------------- # +@dataclass +class PhaseEntry: + """One gated-phase surface and its normalized cardiac stage.""" + + surface: Path + stage: float + + +@dataclass +class SubjectManifest: + """A single subject's training/inference inputs. + + Attributes: + subject_id: Identifier used for output naming. + reference_surface: The subject's SSM reference surface (``.vtp``); the + displacement origin for the Option B convention + (``target = phase.points - reference.points``). + pca_coefficients: JSON file holding the subject's PCA shape-parameter + vector (a flat list of floats). + phases: One :class:`PhaseEntry` per gated phase (at least one). + """ + + subject_id: str + reference_surface: Path + pca_coefficients: Path + phases: list[PhaseEntry] + + +def parse_manifest(manifest_path: Path) -> SubjectManifest: + """Parse a per-subject JSON manifest. + + Paths inside the manifest are resolved relative to the manifest's own + directory unless already absolute. Every phase must declare a ``stage``. + + Args: + manifest_path: Path to the subject manifest JSON file. + + Returns: + The parsed :class:`SubjectManifest`. + + Raises: + FileNotFoundError: If the manifest file does not exist. + ValueError: If required fields are missing, a phase lacks ``stage``, or + no phases are listed. + """ + manifest_path = Path(manifest_path) + if not manifest_path.exists(): + raise FileNotFoundError(f"Manifest not found: {manifest_path}") + + data = json.loads(manifest_path.read_text(encoding="utf-8")) + base = manifest_path.parent + + def _resolve(value: str) -> Path: + p = Path(value) + return p if p.is_absolute() else (base / p) + + for key in ("subject_id", "reference_surface", "pca_coefficients", "phases"): + if key not in data: + raise ValueError(f"Manifest {manifest_path} is missing '{key}'.") + + raw_phases = data["phases"] + if not raw_phases: + raise ValueError(f"Manifest {manifest_path} lists no phases.") + + phases: list[PhaseEntry] = [] + for entry in raw_phases: + if "surface" not in entry or "stage" not in entry: + raise ValueError( + f"Manifest {manifest_path} has a phase missing 'surface' or " + "'stage' (stage must be supplied by the caller)." + ) + phases.append( + PhaseEntry(surface=_resolve(entry["surface"]), stage=float(entry["stage"])) + ) + + return SubjectManifest( + subject_id=str(data["subject_id"]), + reference_surface=_resolve(data["reference_surface"]), + pca_coefficients=_resolve(data["pca_coefficients"]), + phases=phases, + ) + + +def load_pca_coefficients(path: Path) -> np.ndarray: + """Load a PCA shape-parameter vector saved as a JSON list of floats.""" + return np.asarray( + json.loads(Path(path).read_text(encoding="utf-8")), dtype=np.float32 + ) + + +# --------------------------------------------------------------------------- # +# Feature construction (shared by MGN and MLP) # +# --------------------------------------------------------------------------- # +def build_node_features( + mean_coords_norm: np.ndarray, pca_norm: np.ndarray, stage: float +) -> np.ndarray: + """Assemble per-vertex node features ``[coords_norm, pca_norm, stage]``. + + Args: + mean_coords_norm: ``(n_points, 3)`` normalized mean-shape coordinates + (identical for every subject/phase). + pca_norm: ``(n_pca,)`` normalized PCA shape parameters for the subject. + stage: Normalized cardiac stage (RR-interval fraction) for the phase. + + Returns: + ``(n_points, 3 + n_pca + 1)`` float32 feature array. + """ + n = len(mean_coords_norm) + pca_tile = np.tile(pca_norm, (n, 1)) + stage_col = np.full((n, 1), stage, dtype=np.float32) + return np.hstack([mean_coords_norm, pca_tile, stage_col]).astype(np.float32) + + +# --------------------------------------------------------------------------- # +# MGN mesh-graph construction # +# --------------------------------------------------------------------------- # +def mesh_to_edge_index(poly: pv.PolyData) -> "torch.Tensor": + """Build an undirected ``edge_index`` from triangulated PolyData faces. + + Args: + poly: Triangulated surface whose ``faces`` array encodes the topology. + + Returns: + ``(2, n_edges)`` long tensor of undirected edges. + """ + import torch + import torch_geometric.utils as pyg_utils + + # extract_surface may emit quads/other polygons; triangulate (fan) so every + # face is a triangle. vtkTriangleFilter reuses the existing vertices, so the + # point ordering (and thus edge-index correspondence) is preserved. + tri = poly.triangulate() + faces = tri.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 features ``[rel_x, rel_y, rel_z, distance]``.""" + import torch + + ei = edge_index.numpy() + disp = coords[ei[1]] - coords[ei[0]] + dist = np.linalg.norm(disp, axis=1, keepdims=True) + return torch.tensor(np.hstack([disp, dist]), dtype=torch.float32) + + +# --------------------------------------------------------------------------- # +# PCA reconstruction (manifest-free inference) # +# --------------------------------------------------------------------------- # +def reconstruct_reference_points( + mean_mesh: pv.DataSet, pca_model: dict, coeffs: np.ndarray +) -> np.ndarray: + """Reconstruct a subject reference surface from PCA shape parameters. + + Applies the statistical-shape-model equation + ``P = mean + Σ b_i·std_i·eigenvector_i`` on the PCA template *mesh* + (whose ``components`` are defined) and returns the extracted surface points. + Because every subject shares the template topology, the extracted surface + vertex ordering matches the shared mean-shape surface used for training. + + Args: + mean_mesh: PCA template mesh (e.g. ``pca_mean.vtu``) whose point count + matches the model components. + pca_model: Dict with ``eigenvalues`` and ``components`` (the + ``pca_model.json`` format). + coeffs: Subject PCA coefficients ``b_i`` (in units of standard + deviations); shorter/longer than the mode count is truncated. + + Returns: + ``(n_surface_points, 3)`` float32 reconstructed surface points. + + Raises: + ValueError: If the component dimension does not match ``mean_mesh``. + """ + std = np.sqrt(np.asarray(pca_model["eigenvalues"], dtype=np.float64)) + components = np.asarray(pca_model["components"], dtype=np.float64) + expected = mean_mesh.n_points * 3 + if components.shape[1] != expected: + raise ValueError( + f"PCA component dimension {components.shape[1]} does not match " + f"mean mesh ({expected} = 3 x {mean_mesh.n_points} points)." + ) + + b = np.asarray(coeffs, dtype=np.float64) + n_modes = min(len(b), len(std), components.shape[0]) + deform_flat = (b[:n_modes] * std[:n_modes]) @ components[:n_modes] + deform = deform_flat.reshape(-1, 3) + + deformed = mean_mesh.copy(deep=True) + deformed.points = np.asarray(mean_mesh.points, dtype=np.float64) + deform + surface = deformed.extract_surface(algorithm="dataset_surface") + return np.asarray(surface.points, dtype=np.float32) + + +# --------------------------------------------------------------------------- # +# Checkpoint I/O # +# --------------------------------------------------------------------------- # +def uncompiled_state_dict(model: Any) -> dict: + """Return a model's state dict, unwrapping ``torch.compile`` if applied.""" + return cast(dict, getattr(model, "_orig_mod", model).state_dict()) + + +def strip_compile_prefix(state: dict) -> dict: + """Strip the ``_orig_mod.`` prefix that ``torch.compile`` adds to keys.""" + prefix = "_orig_mod." + if any(k.startswith(prefix) for k in state): + return { + k[len(prefix) :] if k.startswith(prefix) else k: v for k, v in state.items() + } + return state + + +# --------------------------------------------------------------------------- # +# Lazy dataset with bounded RAM cache # +# --------------------------------------------------------------------------- # +@dataclass +class _Sample: + """Everything needed to materialize one ``(subject, phase)`` sample.""" + + subject_id: str + pca_norm: np.ndarray # (n_pca,) normalized shape parameters + ref_points: np.ndarray # (n_points, 3) subject reference surface + phase_surface: Path + stage: float + + +class PhaseSampleDataset: + """Lazy provider of ``(node_features, target_displacement)`` samples. + + One item is one ``(subject, phase)`` pair. Node features are rebuilt on + access from the shared normalized mean-shape coordinates plus the subject's + normalized PCA parameters and the phase stage (cheap). Only the phase + surface point arrays are read from disk, and those are held in a bounded + LRU cache so an arbitrarily large training set streams from disk while a + small set stays resident. + + Args: + samples: Flat list of :class:`_Sample` (built by the workflow). + mean_coords_norm: ``(n_points, 3)`` normalized mean-shape coordinates. + displacement_scale: Target normalization factor (targets are divided by + it so displacements land in ``~[-1, 1]``). + cache_max_samples: Maximum decoded phase arrays to cache. ``0`` means + unbounded (all-in-RAM, fastest); a small value forces disk streaming. + """ + + def __init__( + self, + samples: list[_Sample], + mean_coords_norm: np.ndarray, + displacement_scale: float, + cache_max_samples: int = 0, + ) -> None: + self._samples = samples + self._mean_coords_norm = mean_coords_norm.astype(np.float32) + self._displacement_scale = float(displacement_scale) + self._cache_max_samples = int(cache_max_samples) + self._cache: "OrderedDict[Path, np.ndarray]" = OrderedDict() + self._n_points = int(mean_coords_norm.shape[0]) + self._n_features = int(3 + samples[0].pca_norm.shape[0] + 1) if samples else 0 + + def __len__(self) -> int: + return len(self._samples) + + @property + def n_points(self) -> int: + """Vertices per sample (shared across all subjects).""" + return self._n_points + + @property + def n_features(self) -> int: + """Node feature dimension ``3 + n_pca + 1``.""" + return self._n_features + + def _phase_points(self, path: Path) -> np.ndarray: + """Read (and cache) a phase surface's point array.""" + cached = self._cache.get(path) + if cached is not None: + self._cache.move_to_end(path) + return cached + + mesh = pv.read(str(path)) + if mesh.n_points != self._n_points: + raise ValueError( + f"{path} has {mesh.n_points} points, expected {self._n_points}." + ) + points = np.asarray(mesh.points, dtype=np.float32) + + if self._cache_max_samples != 0: + self._cache[path] = points + while len(self._cache) > self._cache_max_samples: + self._cache.popitem(last=False) + return points + + def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray]: + """Return ``(node_features, normalized_target)`` for one sample.""" + sample = self._samples[index] + node_feats = build_node_features( + self._mean_coords_norm, sample.pca_norm, sample.stage + ) + phase_points = self._phase_points(sample.phase_surface) + target = (phase_points - sample.ref_points) / self._displacement_scale + return node_feats, target.astype(np.float32) diff --git a/src/physiotwin4d/segment_chest_total_segmentator.py b/src/physiotwin4d/segment_chest_total_segmentator.py index a8ba499..f2480c3 100644 --- a/src/physiotwin4d/segment_chest_total_segmentator.py +++ b/src/physiotwin4d/segment_chest_total_segmentator.py @@ -209,10 +209,11 @@ def __init__(self, log_level: int | str = logging.INFO): 90: "brain", 15: "esophagus", 16: "trachea", - 133: "body", - 134: "body_trunc", - 135: "body_extremities", - 136: "body_skin", + 133: "body_skin", # 4 in body task + 134: "tissue_subcutaneous_fat", # tissue_4_types + 135: "tissue_torso_fat", + 136: "tissue_skeletal_muscle", + 137: "tissue_intermuscular_fat", }, ), ): @@ -222,15 +223,15 @@ def __init__(self, log_level: int | str = logging.INFO): self._add_extra_taxonomy_groups() self._finalize_other_group() - self.has_highres_heart_license = False + self.has_academic_license = False - def set_has_highres_heart_license(self, has_highres_heart_license: bool) -> None: - """Set whether the highres heart license is available. + def set_has_academic_license(self, has_academic_license: bool) -> None: + """Set whether the academic license is available. Args: - has_highres_heart_license (bool): Whether the highres heart license is available + has_academic_license (bool): Whether the academic license is available """ - self.has_highres_heart_license = has_highres_heart_license + self.has_academic_license = has_academic_license def _add_extra_taxonomy_groups(self) -> None: """Hook for subclasses to add taxonomy groups before finalization. @@ -247,10 +248,12 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: """ Run TotalSegmentator on the preprocessed image and return result. - This implementation runs both the 'total' and 'body' tasks from - TotalSegmentator to ensure comprehensive segmentation. The 'total' task - segments major organs and structures, while the 'body' task provides - body outline segmentation to fill gaps. + This implementation always runs the 'total' task (major organs and + structures). Outside fast mode it also runs the 'lung_vessels' overlay + and the 'body' task; when ``has_academic_license`` is set it additionally + runs the 'heartchambers_highres' and 'tissue_4_types' tasks. The 'body' + task contributes only its skin outline (the skin label) into remaining + background regions; it does not fill gaps with soft tissue. The method uses temporary files for coordinate system conversion between ITK (LPS) and nibabel (RAS) formats, which is required for proper @@ -262,8 +265,8 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: Returns: itk.image: The segmentation labelmap with TotalSegmentator labels. - Background regions from the 'total' task are filled with - soft tissue labels from the 'body' task + The 'body' task's skin label is written into the remaining + background regions as the skin outline. Note: Requires GPU acceleration (device="gpu") for reasonable performance. @@ -303,7 +306,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: final_arr = labelmap_arr_total if not self.fast_mode: - if self.has_highres_heart_license: + if self.has_academic_license: self.log_info("Running heart chambers task") output_nib_image_heart = totalsegmentator( nib_image, @@ -326,6 +329,33 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: # Aorta is not included in heart model. # Should include only a portion of the aorta in the heart model. + self.log_info("Running tissue_4_types task") + output_nib_image_tissue_4_types = totalsegmentator( + nib_image, + task="tissue_4_types", + device="gpu", + nr_thr_resamp=resamp_threads, + ) + labelmap_arr_tissue_4_types = ( + output_nib_image_tissue_4_types.get_fdata().astype(np.uint8) + ) + # 134: "subcutaneous_fat", # tissue_4_types + # 135: "torso_fat", + # 136: "skeletal_muscle", + # 137: "intermuscular_fat" + final_arr = np.where( + labelmap_arr_tissue_4_types == 1, 134, final_arr + ) + final_arr = np.where( + labelmap_arr_tissue_4_types == 2, 135, final_arr + ) + final_arr = np.where( + labelmap_arr_tissue_4_types == 3, 136, final_arr + ) + final_arr = np.where( + labelmap_arr_tissue_4_types == 4, 137, final_arr + ) + self.log_info("Running lung vessels task") output_nib_image_lung = totalsegmentator( nib_image, @@ -352,10 +382,9 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: # Only overwrite the background with body labels mask = final_arr > 0 labelmap_arr_body[mask] = 0 - final_arr = np.where(labelmap_arr_body == 1, 133, final_arr) - final_arr = np.where(labelmap_arr_body == 2, 134, final_arr) - final_arr = np.where(labelmap_arr_body == 3, 135, final_arr) - final_arr = np.where(labelmap_arr_body == 4, 136, final_arr) + final_arr = np.where( + labelmap_arr_body == 4, 133, final_arr + ) # body_skin # To create an ITK image, we save the result and read it back with # ITK. This correctly handles the coordinate system and data @@ -368,7 +397,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: labelmap_arr = itk.array_from_image(labelmap_image).astype(np.uint8) # Add heart around interior regions. - if self.has_highres_heart_license: + if self.has_academic_license: interior_mask = np.isin(labelmap_arr, [141, 142, 143, 144]) # Binarize to foreground value 1 so the dilate/erode calls # below (which use foreground=1) operate on the mask. @@ -387,7 +416,8 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: mask_id = 51 # Heart mask id exterior_arr = exterior_arr * mask_id labelmap_arr = np.where(labelmap_arr == 0, exterior_arr, labelmap_arr) - replace_arr = np.where(labelmap_arr == 133, exterior_arr, 0) + skin_mask = np.isin(labelmap_arr, [133, 134, 135, 136, 137]) + replace_arr = np.where(skin_mask, exterior_arr, 0) labelmap_arr = np.where(replace_arr > 0, exterior_arr, labelmap_arr) labelmap_image = itk.image_from_array(labelmap_arr) diff --git a/src/physiotwin4d/workflow_infer_physicsnemo.py b/src/physiotwin4d/workflow_infer_physicsnemo.py new file mode 100644 index 0000000..15310b0 --- /dev/null +++ b/src/physiotwin4d/workflow_infer_physicsnemo.py @@ -0,0 +1,609 @@ +"""Workflows for inferring cardiac mesh stages with trained PhysicsNeMo models. + +A shared base class :class:`WorkflowInferPhysicsNeMo` loads a model produced by +:mod:`physiotwin4d.workflow_train_physicsnemo` and predicts SSM surface +displacements at requested stages; the concrete +:class:`WorkflowInferPhysicsNeMoMGN` and :class:`WorkflowInferPhysicsNeMoMLP` +subclasses supply the network-specific seams. + +Three prediction entry points share one core displacement predictor: + +- :meth:`WorkflowInferPhysicsNeMo.predict` — manifest-driven prediction of every + gated phase (with error statistics when the phase ground-truth surface exists) + or of arbitrary requested stages. +- :meth:`WorkflowInferPhysicsNeMo.predict_single` — manifest-free single-subject + prediction from a PCA shape-parameter JSON and a target stage, with an + optional ground-truth surface for error reporting. +- :meth:`WorkflowInferPhysicsNeMo.create_deformation_field` — rasterize the + inferred deformation and reference-surface normals onto a caller-supplied + reference image grid. + +PhysicsNeMo (and, for the MGN, PyTorch Geometric) are optional dependencies, +imported lazily so ``import physiotwin4d`` works without them. +""" + +from __future__ import annotations + +import csv +import json +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any, Optional, cast + +import itk +import numpy as np +import pyvista as pv + +from . import physicsnemo_tools as pnt +from .physiotwin4d_base import PhysioTwin4DBase + +if TYPE_CHECKING: # typed for mypy; imported lazily at runtime + import torch + + +class WorkflowInferPhysicsNeMo(PhysioTwin4DBase): + """Base class for inferring cardiac mesh stages from a trained model. + + Not instantiated directly — use :class:`WorkflowInferPhysicsNeMoMGN` or + :class:`WorkflowInferPhysicsNeMoMLP`. Subclasses implement the seams + :meth:`_build_model`, :meth:`_load_extra_artifacts` and + :meth:`_network_predict`, and set the class attribute ``_model_tag``. + """ + + _model_tag: str = "base" + + def __init__( + self, + model_directory: Path, + epoch: Optional[int] = None, + log_level: int | str = logging.INFO, + ) -> None: + """Load a trained model and its normalization statistics. + + Args: + model_directory: Directory written by the matching training + workflow (holds ``_stage_model.pt``, ``pca_mean_surface.vtp`` + and, for the MGN, the shared graph tensors). + epoch: Optional intermittent-checkpoint epoch to load + (``_stage_model_epoch_#####.pt``). When ``None`` the final + weights stored in the main checkpoint are used. + log_level: Logging level. Default: ``logging.INFO``. + + Raises: + FileNotFoundError: If the model checkpoint is missing. + """ + super().__init__(class_name=self.__class__.__name__, log_level=log_level) + import torch + + self.model_directory = Path(model_directory) + checkpoint_file = self.model_directory / f"{self._model_tag}_stage_model.pt" + if not checkpoint_file.exists(): + raise FileNotFoundError(f"Model checkpoint not found: {checkpoint_file}") + + self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.log_info( + "Loading %s model from %s", self._model_tag.upper(), checkpoint_file + ) + meta = torch.load(str(checkpoint_file), map_location="cpu", weights_only=True) + self._meta = meta + + # Normalization statistics. + self.coordinate_mean = np.array(meta["coordinate_mean"], dtype=np.float32) + self.coordinate_scale = np.array(meta["coordinate_scale"], dtype=np.float32) + self.pca_mean = np.array(meta["pca_mean"], dtype=np.float32) + self.pca_scale = np.array(meta["pca_scale"], dtype=np.float32) + self.displacement_scale = float(meta["displacement_scale"]) + + # Shared mean-shape surface (node coordinates + output topology). + mean_surface_file = self.model_directory / "pca_mean_surface.vtp" + if not mean_surface_file.exists(): + raise FileNotFoundError( + f"pca_mean_surface.vtp not found in {self.model_directory}" + ) + self._mean_surface: pv.PolyData = cast( + pv.PolyData, pv.read(str(mean_surface_file)) + ) + self._mean_shape_coords = np.asarray( + self._mean_surface.points, dtype=np.float32 + ) + self._mean_coords_norm = ( + self._mean_shape_coords - self.coordinate_mean + ) / self.coordinate_scale + + # Build the network and load weights. + self._model = self._build_model(meta).to(self._device) + self._load_extra_artifacts() + state = self._load_weights(epoch) + self._model.load_state_dict(pnt.strip_compile_prefix(state)) + self._model.eval() + + # Optional PCA reconstruction assets (manifest-free inference). + self._pca_model: Optional[dict] = None + self._pca_mean_dataset: Optional[pv.DataSet] = None + + # ─────────────────────────── Network seams ───────────────────────────── + def _build_model(self, meta: dict) -> "torch.nn.Module": + """Rebuild the network from checkpoint metadata. Subclass-specific.""" + raise NotImplementedError + + def _load_extra_artifacts(self) -> None: + """Load any architecture-specific artifacts (MGN graph tensors).""" + raise NotImplementedError + + def _network_predict(self, node_feats: np.ndarray) -> np.ndarray: + """Run the network over all vertices; return ``(n, 3)`` raw output.""" + raise NotImplementedError + + # ─────────────────────────── Core predictor ──────────────────────────── + def _load_weights(self, epoch: Optional[int]) -> dict: + """Return the state dict for the requested epoch (or final weights).""" + import torch + + if epoch is None: + return dict(self._meta["model_state_dict"]) + epoch_file = ( + self.model_directory / f"{self._model_tag}_stage_model_epoch_{epoch:05d}.pt" + ) + if not epoch_file.exists(): + raise FileNotFoundError(f"Epoch checkpoint not found: {epoch_file}") + return cast( + dict, torch.load(str(epoch_file), map_location="cpu", weights_only=True) + ) + + def _predict_displacements( + self, pca_coeffs: np.ndarray, stage: float + ) -> np.ndarray: + """Predict per-vertex displacements (mm) for a subject at a stage.""" + pca_norm = (pca_coeffs - self.pca_mean) / self.pca_scale + node_feats = pnt.build_node_features(self._mean_coords_norm, pca_norm, stage) + return self._network_predict(node_feats) * self.displacement_scale + + def _load_pca_assets(self) -> tuple[pv.DataSet, dict]: + """Load (and cache) the PCA template mesh and model for reconstruction.""" + if self._pca_mean_dataset is not None and self._pca_model is not None: + return self._pca_mean_dataset, self._pca_model + + pca_model_file = self.model_directory / "pca_model.json" + if not pca_model_file.exists(): + raise FileNotFoundError( + f"pca_model.json not found in {self.model_directory}; it is " + "required for manifest-free reconstruction. Re-run training with a " + "pca_mean_mesh whose directory contains pca_model.json." + ) + # The PCA template mesh (volume) was copied next to pca_model.json. + mesh_candidates = [ + p + for p in self.model_directory.glob("*") + if p.suffix in (".vtu", ".vtk", ".vtp") and p.name != "pca_mean_surface.vtp" + ] + pca_model = json.loads(pca_model_file.read_text(encoding="utf-8")) + expected = int(np.asarray(pca_model["components"]).shape[1]) // 3 + mesh: Optional[pv.DataSet] = None + for candidate in mesh_candidates: + dataset = pv.read(str(candidate)) + if dataset.n_points == expected: + mesh = dataset + break + if mesh is None: + raise FileNotFoundError( + f"No PCA template mesh with {expected} points found in " + f"{self.model_directory} to match pca_model.json." + ) + self._pca_mean_dataset = mesh + self._pca_model = pca_model + return mesh, pca_model + + # ─────────────────────────── Public API ──────────────────────────────── + def predict( + self, + subject_manifest: Path, + stages: Optional[list[float]] = None, + output_directory: Optional[Path] = None, + ) -> dict[str, Any]: + """Predict a subject's surfaces from a manifest. + + When ``stages`` is ``None`` every gated phase in the manifest is + predicted and, because the phase ground-truth surface is available, + per-phase and per-point error statistics are computed and written. When + ``stages`` is given those arbitrary stages are predicted without error + comparison. + + Args: + subject_manifest: Path to the subject manifest JSON. + stages: Optional list of RR-interval fractions to predict. + output_directory: Output directory; defaults to + ``/``. + + Returns: + Dict with ``subject_id``, ``predicted_surfaces`` (paths), and, in the + phase mode, ``statistics`` and ``rmse_surface``. + """ + manifest = pnt.parse_manifest(subject_manifest) + ref_mesh = cast(pv.PolyData, pv.read(str(manifest.reference_surface))) + ref_points = np.asarray(ref_mesh.points, dtype=np.float32) + pca_coeffs = pnt.load_pca_coefficients(manifest.pca_coefficients) + + out_dir = ( + Path(output_directory) + if output_directory is not None + else self.model_directory / manifest.subject_id + ) + out_dir.mkdir(parents=True, exist_ok=True) + self.log_section("INFER %s [%s]", self._model_tag.upper(), manifest.subject_id) + + if stages is not None: + return self._predict_arbitrary_stages( + manifest.subject_id, ref_mesh, ref_points, pca_coeffs, stages, out_dir + ) + return self._predict_phases(manifest, ref_mesh, ref_points, pca_coeffs, out_dir) + + def _predict_arbitrary_stages( + self, + subject_id: str, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + pca_coeffs: np.ndarray, + stages: list[float], + out_dir: Path, + ) -> dict[str, Any]: + """Predict requested stages without ground-truth comparison.""" + surfaces: list[Path] = [] + for stage in stages: + pred_points = ref_points + self._predict_displacements(pca_coeffs, stage) + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = pred_points + path = ( + out_dir / f"{subject_id}_ssm_surface_pred_s{int(stage * 100):03d}.vtp" + ) + pred_mesh.save(str(path)) + surfaces.append(path) + self.log_info("stage %.3f -> %s", stage, path.name) + return {"subject_id": subject_id, "predicted_surfaces": surfaces} + + def _predict_phases( + self, + manifest: pnt.SubjectManifest, + ref_mesh: pv.PolyData, + ref_points: np.ndarray, + pca_coeffs: np.ndarray, + out_dir: Path, + ) -> dict[str, Any]: + """Predict all manifest phases and compute error statistics vs ground truth.""" + sid = manifest.subject_id + n_points = ref_points.shape[0] + sq_err_sum = np.zeros(n_points, dtype=np.float64) + stats: list[dict] = [] + surfaces: list[Path] = [] + + for phase in manifest.phases: + pred_points = ref_points + self._predict_displacements( + pca_coeffs, phase.stage + ) + pred_mesh = ref_mesh.copy(deep=True) + pred_mesh.points = pred_points + tag = f"s{int(phase.stage * 100):03d}" + path = out_dir / f"{sid}_{tag}_ssm_surface_pred.vtp" + pred_mesh.save(str(path)) + surfaces.append(path) + + actual = np.asarray(pv.read(str(phase.surface)).points, dtype=np.float32) + euclidean = np.linalg.norm(pred_points - actual, axis=1) + sq_err_sum += euclidean.astype(np.float64) ** 2 + stats.append(self._error_row(sid, phase.stage, pred_points, actual)) + self.log_info( + "stage %.3f: mean=%.3f mm max=%.3f mm", + phase.stage, + stats[-1]["mean_error_mm"], + stats[-1]["max_error_mm"], + ) + + point_rmse = np.sqrt(sq_err_sum / len(manifest.phases)).astype(np.float32) + rmse_mesh = ref_mesh.copy(deep=True) + rmse_mesh.point_data["RMSE_mm"] = point_rmse + rmse_file = out_dir / f"{sid}_ssm_surface_rmse.vtp" + rmse_mesh.save(str(rmse_file)) + + stats_file = out_dir / "statistics_per_phase.csv" + with stats_file.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(stats[0].keys())) + writer.writeheader() + writer.writerows(stats) + + return { + "subject_id": sid, + "predicted_surfaces": surfaces, + "rmse_surface": rmse_file, + "statistics": stats, + "statistics_file": stats_file, + } + + def predict_single( + self, + shape_parameters: Path, + stage: float, + ground_truth: Optional[Path] = None, + output_directory: Optional[Path] = None, + ) -> dict[str, Any]: + """Predict one subject at one stage without a manifest. + + The subject reference surface is reconstructed from the PCA shape + parameters (``P = mean + Σ b_i·std_i·eigenvector_i``) in the SSM/PCA + frame, so the prediction is self-consistent even with no reference + surface file. + + Args: + shape_parameters: JSON file with the subject PCA coefficient vector. + stage: Target RR-interval fraction to predict. + ground_truth: Optional surface ``.vtp`` for error reporting. + output_directory: Output directory; defaults to + ``/single_prediction``. + + Returns: + Dict with ``predicted_surface`` (path), ``predicted_points``, and, + when ``ground_truth`` is supplied, ``statistics``. + """ + coeffs = pnt.load_pca_coefficients(shape_parameters) + mean_mesh, pca_model = self._load_pca_assets() + ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) + pred_points = ref_points + self._predict_displacements(coeffs, stage) + + out_dir = ( + Path(output_directory) + if output_directory is not None + else self.model_directory / "single_prediction" + ) + out_dir.mkdir(parents=True, exist_ok=True) + + pred_mesh = self._mean_surface.copy(deep=True) + pred_mesh.points = pred_points + stem = Path(shape_parameters).stem + path = out_dir / f"{stem}_pred_s{int(stage * 100):03d}.vtp" + pred_mesh.save(str(path)) + self.log_info("single prediction stage %.3f -> %s", stage, path.name) + + result: dict[str, Any] = { + "predicted_surface": path, + "predicted_points": pred_points, + } + if ground_truth is not None: + actual = np.asarray(pv.read(str(ground_truth)).points, dtype=np.float32) + result["statistics"] = self._error_row(stem, stage, pred_points, actual) + self.log_info( + "ground-truth error: mean=%.3f mm max=%.3f mm", + result["statistics"]["mean_error_mm"], + result["statistics"]["max_error_mm"], + ) + return result + + def create_deformation_field( + self, + shape_parameters: Path, + stage: float, + reference_image: itk.Image, + output_directory: Optional[Path] = None, + ) -> dict[str, Any]: + """Rasterize the inferred deformation onto a reference image grid. + + Each mesh vertex is binned by its **reference (undeformed) position** + into ``reference_image``'s voxel grid. Each voxel of the deformation + field holds the mean network displacement ``(dx, dy, dz)`` of the + vertices that fall in it; each voxel of the normal image holds the mean + (renormalized) reference-surface normal of those vertices. Empty voxels + are zero. + + Args: + shape_parameters: JSON file with the subject PCA coefficient vector. + stage: Target RR-interval fraction for the deformation. + reference_image: The frame's image; defines the output grid geometry + (size, spacing, origin, direction). + output_directory: If given, the two images are written there as + compressed ``.mha`` files. + + Returns: + Dict with ``deformation_field`` and ``normal_image`` (ITK vector + images) and, when written, their paths. + """ + coeffs = pnt.load_pca_coefficients(shape_parameters) + mean_mesh, pca_model = self._load_pca_assets() + ref_points = pnt.reconstruct_reference_points(mean_mesh, pca_model, coeffs) + disps = self._predict_displacements(coeffs, stage) + + # Reference (undeformed) surface normals. + ref_surface = self._mean_surface.copy(deep=True) + ref_surface.points = ref_points + ref_surface = ref_surface.compute_normals( + point_normals=True, cell_normals=False, auto_orient_normals=True + ) + normals = np.asarray(ref_surface.point_data["Normals"], dtype=np.float64) + + size = itk.size(reference_image) # x, y, z + sx, sy, sz = int(size[0]), int(size[1]), int(size[2]) + disp_sum = np.zeros((sz, sy, sx, 3), dtype=np.float64) + normal_sum = np.zeros((sz, sy, sx, 3), dtype=np.float64) + count = np.zeros((sz, sy, sx), dtype=np.float64) + + for i in range(ref_points.shape[0]): + point = [float(c) for c in ref_points[i]] + idx = reference_image.TransformPhysicalPointToIndex(point) + ix, iy, iz = int(idx[0]), int(idx[1]), int(idx[2]) + if 0 <= ix < sx and 0 <= iy < sy and 0 <= iz < sz: + disp_sum[iz, iy, ix] += disps[i] + normal_sum[iz, iy, ix] += normals[i] + count[iz, iy, ix] += 1.0 + + occupied = count > 0 + disp_field = np.zeros_like(disp_sum, dtype=np.float32) + normal_field = np.zeros_like(normal_sum, dtype=np.float32) + disp_field[occupied] = (disp_sum[occupied] / count[occupied, None]).astype( + np.float32 + ) + mean_normal = normal_sum[occupied] / count[occupied, None] + norm = np.linalg.norm(mean_normal, axis=1, keepdims=True) + norm = np.where(norm == 0.0, 1.0, norm) + normal_field[occupied] = (mean_normal / norm).astype(np.float32) + + deformation_image = self._vector_image_like(disp_field, reference_image) + normal_image = self._vector_image_like(normal_field, reference_image) + self.log_info( + "Deformation field: %d/%d voxels populated by %d vertices", + int(occupied.sum()), + sx * sy * sz, + ref_points.shape[0], + ) + + result: dict[str, Any] = { + "deformation_field": deformation_image, + "normal_image": normal_image, + } + if output_directory is not None: + out_dir = Path(output_directory) + out_dir.mkdir(parents=True, exist_ok=True) + field_path = out_dir / "deformation_field.mha" + normal_path = out_dir / "surface_normal_field.mha" + itk.imwrite(deformation_image, str(field_path), compression=True) + itk.imwrite(normal_image, str(normal_path), compression=True) + result["deformation_field_file"] = field_path + result["normal_image_file"] = normal_path + return result + + @staticmethod + def _vector_image_like(array: np.ndarray, reference_image: itk.Image) -> itk.Image: + """Wrap a ``(z, y, x, 3)`` array as an ITK vector image on ``reference``'s grid.""" + image = itk.image_from_array(np.ascontiguousarray(array), is_vector=True) + image.SetSpacing(reference_image.GetSpacing()) + image.SetOrigin(reference_image.GetOrigin()) + image.SetDirection(reference_image.GetDirection()) + return image + + @staticmethod + def _error_row( + subject_id: str, stage: float, pred: np.ndarray, actual: np.ndarray + ) -> dict: + """Per-phase error statistics between predicted and actual points.""" + errors = pred - actual + euclidean = np.linalg.norm(errors, axis=1) + return { + "subject_id": subject_id, + "stage": stage, + "n_points": int(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()), + } + + +class WorkflowInferPhysicsNeMoMGN(WorkflowInferPhysicsNeMo): + """Infer cardiac mesh stages with a trained PhysicsNeMo MeshGraphNet.""" + + _model_tag = "mgn" + + def _build_model(self, meta: dict) -> "torch.nn.Module": + + try: + import torch_geometric # noqa: F401 - needed by the graph seams + + from physicsnemo.models.meshgraphnet import MeshGraphNet + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "The MGN inferencer requires PhysicsNeMo and PyTorch Geometric. " + 'Install with: pip install "physiotwin4d[physicsnemo]" && ' + "pip install torch-geometric" + ) from exc + + num_layers = int(meta.get("num_layers", 2)) + hidden_dim = int(meta["hidden_dim"]) + model = MeshGraphNet( + input_dim_nodes=int(meta["in_features"]), + input_dim_edges=int(meta.get("input_dim_edges", 4)), + output_dim=3, + processor_size=int(meta["processor_size"]), + hidden_dim_processor=hidden_dim, + hidden_dim_node_encoder=hidden_dim, + num_layers_node_encoder=num_layers, + hidden_dim_node_decoder=hidden_dim, + num_layers_node_decoder=num_layers, + hidden_dim_edge_encoder=hidden_dim, + num_layers_edge_encoder=num_layers, + num_layers_edge_processor=num_layers, + num_layers_node_processor=num_layers, + aggregation="mean", + num_processor_checkpoint_segments=int( + meta.get("num_processor_checkpoint_segments", 0) + ), + ) + return cast("torch.nn.Module", model) + + def _load_extra_artifacts(self) -> None: + import torch + from torch_geometric.data import Data + + edge_index = torch.load( + str(self.model_directory / "shared_edge_index.pt"), + map_location="cpu", + weights_only=True, + ) + edge_feats = torch.load( + str(self.model_directory / "shared_edge_features.pt"), + map_location="cpu", + weights_only=True, + ) + self._shared_edge_index = edge_index + self._shared_edge_feats = edge_feats.to(self._device) + self._shared_graph = Data( + edge_index=edge_index, num_nodes=len(self._mean_shape_coords) + ).to(self._device) + + def _network_predict(self, node_feats: np.ndarray) -> np.ndarray: + import torch + + nf = torch.from_numpy(node_feats.astype(np.float32)).to(self._device) + with torch.no_grad(): + pred = self._model(nf, self._shared_edge_feats, self._shared_graph) + return np.asarray(pred.cpu().numpy(), dtype=np.float32) + + +class WorkflowInferPhysicsNeMoMLP(WorkflowInferPhysicsNeMo): + """Infer cardiac mesh stages with a trained PhysicsNeMo FullyConnected model.""" + + _model_tag = "mlp" + _INFER_CHUNK = 262144 + + def _build_model(self, meta: dict) -> "torch.nn.Module": + + try: + from physicsnemo.models.mlp import FullyConnected + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "The MLP inferencer requires PhysicsNeMo, an optional dependency. " + 'Install with: pip install "physiotwin4d[physicsnemo]"' + ) from exc + + 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, + ) + return cast("torch.nn.Module", model) + + def _load_extra_artifacts(self) -> None: + # MLP has no shared graph artifacts. + return None + + def _network_predict(self, node_feats: np.ndarray) -> np.ndarray: + import torch + + chunks: list[np.ndarray] = [] + with torch.no_grad(): + for start in range(0, len(node_feats), self._INFER_CHUNK): + block = node_feats[start : start + self._INFER_CHUNK].astype(np.float32) + tensor = torch.from_numpy(block).to(self._device) + chunks.append(self._model(tensor).cpu().numpy()) + return np.vstack(chunks).astype(np.float32) diff --git a/src/physiotwin4d/workflow_train_physicsnemo.py b/src/physiotwin4d/workflow_train_physicsnemo.py new file mode 100644 index 0000000..430f5ac --- /dev/null +++ b/src/physiotwin4d/workflow_train_physicsnemo.py @@ -0,0 +1,839 @@ +"""Workflows for training PhysicsNeMo cardiac mesh-stage models. + +A shared base class :class:`WorkflowTrainPhysicsNeMo` holds every step common to +the two supported networks; the concrete +:class:`WorkflowTrainPhysicsNeMoMGN` (MeshGraphNet) and +:class:`WorkflowTrainPhysicsNeMoMLP` (fully connected) subclasses supply only +the network-specific seams. Both learn the same task as tutorials 9a/9b: given +per-vertex features ``[mean_shape_x, mean_shape_y, mean_shape_z, pca_c1..cN, +stage]`` predict a per-vertex displacement ``(dx, dy, dz)`` from the subject's +SSM reference surface (the Option B convention). + +Design highlights: + +- **Data is a list of per-subject manifest files** (see + :func:`physiotwin4d.physicsnemo_tools.parse_manifest`). The caller chooses the + train / validation / held-out-test split externally; the trainer receives the + training manifests and validation manifest(s) and reports validation RMSE + intermittently as training proceeds. +- **The dataset streams lazily** through + :class:`physiotwin4d.physicsnemo_tools.PhaseSampleDataset` with a bounded RAM + cache, so the training set need not fit in memory (unlike the tutorials, which + pre-stack the whole dataset onto the GPU). +- **Coordinates are always the PCA mean-shape surface** (shared across + subjects); the subject is described by its PCA parameters and the stage. + +MLP note: the tutorial MLP drew mini-batches of individual points shuffled +across all subjects/phases. Streaming by file makes the natural batch unit a +group of whole ``(subject, phase)`` samples, so the MLP batches several samples +per step and shuffles points *within* the batch to retain gradient mixing +(``batch_size`` therefore counts samples, not points). The MGN keeps per-sample +vertex order intact because it indexes the shared mesh graph. + +PhysicsNeMo (and, for the MGN, PyTorch Geometric) are optional dependencies; +they are imported lazily so ``import physiotwin4d`` works without them. +""" + +from __future__ import annotations + +import csv +import json +import logging +import shutil +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Any, Optional, cast + +import numpy as np +import pyvista as pv + +from . import physicsnemo_tools as pnt +from .physicsnemo_tools import PhaseSampleDataset, SubjectManifest, _Sample +from .physiotwin4d_base import PhysioTwin4DBase +from .test_tools import TestTools + +if TYPE_CHECKING: # typed for mypy; imported lazily at runtime + import torch + + +class WorkflowTrainPhysicsNeMo(PhysioTwin4DBase): + """Base class for training a PhysicsNeMo cardiac mesh-stage model. + + Not instantiated directly — use :class:`WorkflowTrainPhysicsNeMoMGN` or + :class:`WorkflowTrainPhysicsNeMoMLP`. Subclasses implement the network + seams (:meth:`_build_model`, :meth:`_setup_model_inputs`, :meth:`_forward`, + :meth:`_checkpoint_extra`, :meth:`_save_extra_artifacts`) and set the class + attributes ``_model_tag``, ``_architecture_name`` and + ``_shuffle_points_within_batch``. + """ + + # Network identity / behavior — overridden by subclasses. + _model_tag: str = "base" + _architecture_name: str = "base" + _shuffle_points_within_batch: bool = False + + def __init__( + self, + train_manifests: list[Path], + val_manifests: list[Path], + pca_mean_mesh: Path, + output_directory: Path, + resume_from: Optional[Path] = None, + log_level: int | str = logging.INFO, + ) -> None: + """Initialize the trainer. + + Args: + train_manifests: Per-subject manifest files for the training set. + val_manifests: Per-subject manifest files for the validation set + (used for intermittent RMSE reporting during training). May be + empty to skip validation. + pca_mean_mesh: PCA template mesh whose point count matches + ``pca_model.json`` (typically ``pca_mean.vtu``). Its extracted + surface defines the shared node coordinates and — for the MGN — + the mesh-graph topology. The sibling ``pca_model.json`` (if + present) is copied into ``output_directory`` for inference. + output_directory: Directory for checkpoints, metadata and logs. + resume_from: Optional ``*_stage_model.pt`` to resume from; its + normalization statistics are inherited so the loaded weights stay + valid, and a fresh numbered output directory is used. + log_level: Logging level. Default: ``logging.INFO``. + + Raises: + ValueError: If ``train_manifests`` is empty. + FileNotFoundError: If ``pca_mean_mesh`` does not exist. + """ + super().__init__(class_name=self.__class__.__name__, log_level=log_level) + + if not train_manifests: + raise ValueError("train_manifests cannot be empty.") + pca_mean_mesh = Path(pca_mean_mesh) + if not pca_mean_mesh.exists(): + raise FileNotFoundError(f"pca_mean_mesh not found: {pca_mean_mesh}") + + self.train_manifest_paths = [Path(p) for p in train_manifests] + self.val_manifest_paths = [Path(p) for p in val_manifests] + self.pca_mean_mesh = pca_mean_mesh + self.output_directory = Path(output_directory) + self.resume_from = Path(resume_from) if resume_from is not None else None + + # PCA assets shared by every subject. + self._pca_mean_dataset: pv.DataSet = pv.read(str(pca_mean_mesh)) + self._mean_surface: pv.PolyData = self._pca_mean_dataset.extract_surface( + algorithm="dataset_surface" + ) + self._mean_shape_coords = np.asarray( + self._mean_surface.points, dtype=np.float32 + ) + self._pca_model_path: Optional[Path] = None + candidate = pca_mean_mesh.parent / "pca_model.json" + if candidate.exists(): + self._pca_model_path = candidate + + # Shared hyper-parameters (subclasses may override the batch default). + self.epochs: int = 1500 + self.batch_size: int = 4 + self.learning_rate: float = 1.0e-3 + self.cache_max_samples: int = 0 + self.rmse_log_interval: int = 100 + self.loss_log_interval: int = 10 + self.seed: int = 42 + + # Results (populated by process()). + self.checkpoint_file: Optional[Path] = None + self.metadata_file: Optional[Path] = None + self.training_loss: Optional[list[float]] = None + self.val_rmse_log: Optional[list[dict]] = None + + # ─────────────────────────── Tuning setters ──────────────────────────── + def set_epochs(self, epochs: int) -> None: + """Set the number of training epochs.""" + if epochs < 1: + raise ValueError(f"epochs must be >= 1, got {epochs}") + self.epochs = epochs + + def set_batch_size(self, batch_size: int) -> None: + """Set the mini-batch size, measured in ``(subject, phase)`` samples.""" + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}") + self.batch_size = batch_size + + def set_learning_rate(self, learning_rate: float) -> None: + """Set the Adam learning rate.""" + if learning_rate <= 0.0: + raise ValueError(f"learning_rate must be > 0, got {learning_rate}") + self.learning_rate = learning_rate + + def set_cache_size(self, cache_max_samples: int) -> None: + """Set the RAM cache budget (decoded phase arrays); ``0`` = unbounded.""" + if cache_max_samples < 0: + raise ValueError(f"cache_max_samples must be >= 0, got {cache_max_samples}") + self.cache_max_samples = cache_max_samples + + # ─────────────────────────── Network seams ───────────────────────────── + def _build_model(self, in_features: int) -> "torch.nn.Module": + """Construct the (uncompiled) network. Implemented by subclasses.""" + raise NotImplementedError + + def _setup_model_inputs(self, device: "torch.device") -> None: + """Prepare any shared per-forward inputs (MGN graph tensors).""" + raise NotImplementedError + + def _forward( + self, model: "torch.nn.Module", node_feats: "torch.Tensor", batch_len: int + ) -> "torch.Tensor": + """Run the network for a flattened ``(batch_len * n_points, F)`` batch.""" + raise NotImplementedError + + def _checkpoint_extra(self) -> dict: + """Return architecture-specific fields to store in the checkpoint.""" + raise NotImplementedError + + def _save_extra_artifacts(self, output_dir: Path) -> None: + """Save any architecture-specific artifacts (MGN graph tensors).""" + raise NotImplementedError + + # ─────────────────────────── Main workflow ───────────────────────────── + def process(self) -> dict[str, Any]: + """Train the model and write checkpoints, metadata and logs. + + Returns: + Dict with ``output_directory``, ``checkpoint``, ``metadata``, + ``training_loss`` and ``val_rmse_log``. + """ + import torch + + self.log_section( + "STARTING PHYSICSNEMO %s TRAINING WORKFLOW", self._model_tag.upper() + ) + + epochs = self.epochs + if TestTools.running_as_test(): + epochs = min(epochs, 2) + self.log_info("Test mode: reducing epochs to %d", epochs) + + output_dir = self._resolve_output_dir() + output_dir.mkdir(parents=True, exist_ok=True) + self.log_info("Output directory: %s", output_dir) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.log_info("Device: %s", device.type) + + subjects = self._load_subjects() + resume_ckpt = self._load_resume_checkpoint() + stats = self._compute_normalization(subjects, resume_ckpt) + + train_dataset, val_dataset = self._build_datasets(subjects, stats) + self.log_info( + "Training samples: %d, validation samples: %d, in_features=%d, " + "displacement_scale=%.4f mm", + len(train_dataset), + len(val_dataset), + train_dataset.n_features, + stats["displacement_scale"], + ) + + model, losses, rmse_log = self._train( + train_dataset, val_dataset, stats, device, epochs, output_dir + ) + self._save_model(model, subjects, stats, losses, rmse_log, output_dir, epochs) + + self.log_section("PHYSICSNEMO %s TRAINING COMPLETE", self._model_tag.upper()) + return { + "output_directory": output_dir, + "checkpoint": self.checkpoint_file, + "metadata": self.metadata_file, + "training_loss": losses, + "val_rmse_log": rmse_log, + } + + # ─────────────────────────── Internal steps ──────────────────────────── + def _resolve_output_dir(self) -> Path: + """Return the output directory, using a fresh sibling when resuming.""" + base = self.output_directory + if self.resume_from is None or 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 _load_subjects(self) -> dict[str, dict]: + """Parse every manifest and load reference points + PCA coefficients.""" + n_points = len(self._mean_shape_coords) + subjects: dict[str, dict] = {} + + def _load(paths: list[Path], split: str) -> None: + for manifest_path in paths: + manifest: SubjectManifest = pnt.parse_manifest(manifest_path) + if manifest.subject_id in subjects: + raise ValueError( + f"Duplicate subject_id '{manifest.subject_id}': already " + f"loaded in the '{subjects[manifest.subject_id]['split']}' " + f"split, seen again in the '{split}' split. Each subject " + "must appear in exactly one manifest." + ) + ref_mesh = pv.read(str(manifest.reference_surface)) + ref_points = np.asarray(ref_mesh.points, dtype=np.float32) + if ref_points.shape[0] != n_points: + raise ValueError( + f"{manifest.reference_surface} has {ref_points.shape[0]} " + f"points, expected {n_points} (mean-shape topology)." + ) + subjects[manifest.subject_id] = { + "split": split, + "ref_points": ref_points, + "pca_coeffs": pnt.load_pca_coefficients(manifest.pca_coefficients), + "phases": manifest.phases, + } + + _load(self.train_manifest_paths, "train") + _load(self.val_manifest_paths, "val") + n_train = sum(1 for s in subjects.values() if s["split"] == "train") + if n_train == 0: + raise ValueError("No training subjects were loaded.") + return subjects + + def _load_resume_checkpoint(self) -> Optional[dict]: + """Load prior-run normalization statistics when resuming.""" + if self.resume_from is None: + return None + import torch + + self.log_info("Resuming from %s", self.resume_from) + return cast( + dict, + torch.load(str(self.resume_from), map_location="cpu", weights_only=True), + ) + + def _compute_normalization( + self, subjects: dict[str, dict], resume_ckpt: Optional[dict] + ) -> dict: + """Compute (or inherit) coordinate, PCA and displacement statistics.""" + if resume_ckpt is not None: + return { + "coordinate_mean": np.array(resume_ckpt["coordinate_mean"], np.float32), + "coordinate_scale": np.array( + resume_ckpt["coordinate_scale"], np.float32 + ), + "pca_mean": np.array(resume_ckpt["pca_mean"], np.float32), + "pca_scale": np.array(resume_ckpt["pca_scale"], np.float32), + "displacement_scale": float(resume_ckpt["displacement_scale"]), + } + + coord = self._mean_shape_coords + coordinate_mean = coord.mean(axis=0) + coordinate_scale = np.where(coord.std(axis=0) == 0.0, 1.0, coord.std(axis=0)) + + train_pca = np.vstack( + [s["pca_coeffs"] for s in subjects.values() if s["split"] == "train"] + ) + pca_mean = train_pca.mean(axis=0) + pca_scale = np.where(train_pca.std(axis=0) == 0.0, 1.0, train_pca.std(axis=0)) + + displacement_scale = self._compute_displacement_scale(subjects) + return { + "coordinate_mean": coordinate_mean.astype(np.float32), + "coordinate_scale": coordinate_scale.astype(np.float32), + "pca_mean": pca_mean.astype(np.float32), + "pca_scale": pca_scale.astype(np.float32), + "displacement_scale": displacement_scale, + } + + def _compute_displacement_scale(self, subjects: dict[str, dict]) -> float: + """One streaming pass over training targets for the max abs displacement.""" + n_points = len(self._mean_shape_coords) + max_abs = 0.0 + for data in subjects.values(): + if data["split"] != "train": + continue + ref_points = data["ref_points"] + for phase in data["phases"]: + mesh = pv.read(str(phase.surface)) + if mesh.n_points != n_points: + raise ValueError( + f"{phase.surface} has {mesh.n_points} points, " + f"expected {n_points}." + ) + disp = np.asarray(mesh.points, dtype=np.float32) - ref_points + max_abs = max(max_abs, float(np.max(np.abs(disp)))) + return max_abs if max_abs > 0.0 else 1.0 + + def _build_datasets( + self, subjects: dict[str, dict], stats: dict + ) -> tuple[PhaseSampleDataset, PhaseSampleDataset]: + """Build lazy train and validation datasets sharing the mean coords.""" + mean_coords_norm = (self._mean_shape_coords - stats["coordinate_mean"]) / stats[ + "coordinate_scale" + ] + + def _samples(split: str) -> list[_Sample]: + out: list[_Sample] = [] + for sid, data in sorted(subjects.items()): + if data["split"] != split: + continue + pca_norm = (data["pca_coeffs"] - stats["pca_mean"]) / stats["pca_scale"] + for phase in data["phases"]: + out.append( + _Sample( + subject_id=sid, + pca_norm=pca_norm.astype(np.float32), + ref_points=data["ref_points"], + phase_surface=phase.surface, + stage=phase.stage, + ) + ) + return out + + train_dataset = PhaseSampleDataset( + _samples("train"), + mean_coords_norm, + stats["displacement_scale"], + self.cache_max_samples, + ) + val_dataset = PhaseSampleDataset( + _samples("val"), + mean_coords_norm, + stats["displacement_scale"], + self.cache_max_samples, + ) + return train_dataset, val_dataset + + def _iter_batches( + self, dataset: PhaseSampleDataset, rng: Any, shuffle: bool + ) -> Any: + """Yield ``(node_feats, targets, batch_len)`` flattened mini-batches.""" + n = len(dataset) + order = rng.permutation(n) if shuffle else np.arange(n) + for start in range(0, n, self.batch_size): + idx = order[start : start + self.batch_size] + pairs = [dataset[int(i)] for i in idx] + node_feats = np.vstack([p[0] for p in pairs]) + targets = np.vstack([p[1] for p in pairs]) + if shuffle and self._shuffle_points_within_batch: + perm = rng.permutation(len(node_feats)) + node_feats = node_feats[perm] + targets = targets[perm] + yield node_feats, targets, len(idx) + + def _train( + self, + train_dataset: PhaseSampleDataset, + val_dataset: PhaseSampleDataset, + stats: dict, + device: "torch.device", + epochs: int, + output_dir: Path, + ) -> tuple["torch.nn.Module", list[float], list[dict]]: + """Run the training loop, returning the model and loss/RMSE logs.""" + import torch + + torch.manual_seed(self.seed) + rng = np.random.default_rng(self.seed) + in_features = train_dataset.n_features + + model = self._build_model(in_features).to(device) + if self.resume_from is not None: + ckpt = torch.load( + str(self.resume_from), map_location=device, weights_only=True + ) + state = ckpt.get("model_state_dict", ckpt) + model.load_state_dict(pnt.strip_compile_prefix(state)) + self.log_info("Loaded model weights from %s", self.resume_from) + + self._setup_model_inputs(device) + + if sys.platform != "win32": + try: + model = cast("torch.nn.Module", torch.compile(model)) + self.log_info("torch.compile enabled.") + except Exception as exc: # pragma: no cover - platform dependent + self.log_info("torch.compile skipped (%s).", exc) + else: + self.log_info("torch.compile skipped on Windows.") + + optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate) + loss_fn = torch.nn.MSELoss() + displacement_scale = stats["displacement_scale"] + + losses: list[float] = [] + rmse_log: list[dict] = [] + for epoch in range(epochs): + model.train() + epoch_loss = 0.0 + n_rows = 0 + for node_feats, targets, batch_len in self._iter_batches( + train_dataset, rng, shuffle=True + ): + nf = torch.from_numpy(node_feats).to(device) + tgt = torch.from_numpy(targets).to(device) + optimizer.zero_grad(set_to_none=True) + with self._autocast(device): + pred = self._forward(model, nf, batch_len) + loss = loss_fn(pred, tgt) + loss.backward() + optimizer.step() + epoch_loss += float(loss.detach()) * len(nf) + n_rows += len(nf) + losses.append(epoch_loss / max(n_rows, 1)) + + if (epoch + 1) % self.loss_log_interval == 0 or epoch + 1 == epochs: + self.log_info( + " epoch %05d/%d loss=%.6f", epoch + 1, epochs, losses[-1] + ) + + if (epoch + 1) % self.rmse_log_interval == 0 or epoch + 1 == epochs: + train_rmse = self._evaluate_rmse( + model, train_dataset, displacement_scale, device + ) + val_rmse = ( + self._evaluate_rmse(model, val_dataset, displacement_scale, device) + if len(val_dataset) > 0 + else float("nan") + ) + rmse_log.append( + { + "epoch": epoch + 1, + "train_rmse_mm": train_rmse, + "val_rmse_mm": val_rmse, + } + ) + ckpt_path = ( + output_dir + / f"{self._model_tag}_stage_model_epoch_{epoch + 1:05d}.pt" + ) + torch.save(pnt.uncompiled_state_dict(model), ckpt_path) + self.log_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() + return model, losses, rmse_log + + def _autocast(self, device: "torch.device") -> Any: + """BF16 autocast on CUDA; a no-op context elsewhere.""" + import contextlib + + import torch + + if device.type == "cuda": + return torch.amp.autocast(device.type, dtype=torch.bfloat16) + return contextlib.nullcontext() + + def _evaluate_rmse( + self, + model: "torch.nn.Module", + dataset: PhaseSampleDataset, + displacement_scale: float, + device: "torch.device", + ) -> float: + """Euclidean per-point RMSE in mm over a dataset.""" + import torch + + rng = np.random.default_rng(0) + model.eval() + total_sq = 0.0 + n_points = 0 + with torch.no_grad(): + for node_feats, targets, batch_len in self._iter_batches( + dataset, rng, shuffle=False + ): + nf = torch.from_numpy(node_feats).to(device) + pred = self._forward(model, nf, batch_len).cpu().numpy() + err = (pred - targets) * displacement_scale + total_sq += float(np.sum(err**2)) + n_points += err.shape[0] + model.train() + return float(np.sqrt(total_sq / max(n_points, 1))) + + def _save_model( + self, + model: "torch.nn.Module", + subjects: dict[str, dict], + stats: dict, + losses: list[float], + rmse_log: list[dict], + output_dir: Path, + epochs: int, + ) -> None: + """Persist the checkpoint, metadata, logs and shared PCA assets.""" + import torch + + in_features = 3 + int(stats["pca_mean"].shape[0]) + 1 + tag = self._model_tag + checkpoint_file = output_dir / f"{tag}_stage_model.pt" + metadata_file = output_dir / f"{tag}_stage_model_metadata.json" + + train_ids = sorted(s for s, d in subjects.items() if d["split"] == "train") + val_ids = sorted(s for s, d in subjects.items() if d["split"] == "val") + + checkpoint: dict[str, Any] = { + "model_state_dict": pnt.uncompiled_state_dict(model), + "architecture": self._architecture_name, + "in_features": in_features, + "n_pca": int(stats["pca_mean"].shape[0]), + "coordinate_mean": stats["coordinate_mean"].tolist(), + "coordinate_scale": stats["coordinate_scale"].tolist(), + "pca_mean": stats["pca_mean"].tolist(), + "pca_scale": stats["pca_scale"].tolist(), + "displacement_scale": stats["displacement_scale"], + "train_subject_ids": train_ids, + "val_subject_ids": val_ids, + "resumed_from": str(self.resume_from) if self.resume_from else None, + } + checkpoint.update(self._checkpoint_extra()) + torch.save(checkpoint, checkpoint_file) + + n_pca = int(stats["pca_mean"].shape[0]) + input_feature_names = ( + ["mean_shape_x", "mean_shape_y", "mean_shape_z"] + + [f"pca_c{i + 1}" for i in range(n_pca)] + + ["stage"] + ) + metadata = { + "architecture": self._architecture_name, + "input_features": input_feature_names, + "output_features": ["dx", "dy", "dz"], + "in_features": in_features, + "n_mesh_points": int(self._mean_shape_coords.shape[0]), + "epochs": epochs, + "learning_rate": self.learning_rate, + "batch_size_samples": self.batch_size, + "coordinate_mean": stats["coordinate_mean"].tolist(), + "coordinate_scale": stats["coordinate_scale"].tolist(), + "pca_mean": stats["pca_mean"].tolist(), + "pca_scale": stats["pca_scale"].tolist(), + "displacement_scale": stats["displacement_scale"], + "displacement_convention": "Option B: relative to subject reference surface", + "resumed_from": str(self.resume_from) if self.resume_from else None, + } + metadata.update(self._checkpoint_extra()) + metadata_file.write_text(json.dumps(metadata, indent=2), encoding="utf-8") + + (output_dir / "training_losses.json").write_text( + json.dumps(losses, indent=2), encoding="utf-8" + ) + (output_dir / "training_validation_rmse.json").write_text( + json.dumps(rmse_log, indent=2), encoding="utf-8" + ) + with (output_dir / "training_validation_rmse.csv").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 model directory is self-contained for inference. + shutil.copy2(self.pca_mean_mesh, output_dir / self.pca_mean_mesh.name) + self._mean_surface.save(str(output_dir / "pca_mean_surface.vtp")) + if self._pca_model_path is not None: + shutil.copy2(self._pca_model_path, output_dir / "pca_model.json") + + self._save_extra_artifacts(output_dir) + + self.checkpoint_file = checkpoint_file + self.metadata_file = metadata_file + self.training_loss = losses + self.val_rmse_log = rmse_log + self.log_info("Model saved to %s", checkpoint_file) + + +class WorkflowTrainPhysicsNeMoMGN(WorkflowTrainPhysicsNeMo): + """Train a PhysicsNeMo :class:`MeshGraphNet` on cardiac mesh stages. + + The mesh-graph topology and edge features are extracted once from the shared + mean-shape surface and reused for every ``(subject, phase)`` sample; PyTorch + Geometric batches join disconnected sub-graphs. + """ + + _model_tag = "mgn" + _architecture_name = "physicsnemo.models.meshgraphnet.MeshGraphNet" + _shuffle_points_within_batch = False + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.batch_size = 4 # graphs per step + self.processor_size: int = 3 + self.hidden_dim: int = 128 + self.num_layers: int = 2 + self.num_processor_checkpoint_segments: int = 0 + # Runtime MGN state (set in _setup_model_inputs). + self._device: Optional["torch.device"] = None + self._shared_graph: Any = None + self._shared_edge_index: Any = None + self._shared_edge_feats: Any = None + self._batched_graph_cache: dict[int, tuple[Any, Any]] = {} + + def set_processor_size(self, processor_size: int) -> None: + """Set the number of message-passing hops.""" + if processor_size < 1: + raise ValueError(f"processor_size must be >= 1, got {processor_size}") + self.processor_size = processor_size + + def set_hidden_dim(self, hidden_dim: int) -> None: + """Set the processor/encoder/decoder hidden dimension.""" + if hidden_dim < 1: + raise ValueError(f"hidden_dim must be >= 1, got {hidden_dim}") + self.hidden_dim = hidden_dim + + def set_num_layers(self, num_layers: int) -> None: + """Set the MLP layer count inside each encoder/processor/decoder block.""" + if num_layers < 1: + raise ValueError(f"num_layers must be >= 1, got {num_layers}") + self.num_layers = num_layers + + def _build_model(self, in_features: int) -> "torch.nn.Module": + + try: + import torch_geometric # noqa: F401 - needed by the graph seams + + from physicsnemo.models.meshgraphnet import MeshGraphNet + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "The MGN trainer requires PhysicsNeMo and PyTorch Geometric. " + 'Install with: pip install "physiotwin4d[physicsnemo]" && ' + "pip install torch-geometric" + ) from exc + + model = MeshGraphNet( + input_dim_nodes=in_features, + input_dim_edges=4, # rel_x, rel_y, rel_z, distance + output_dim=3, + processor_size=self.processor_size, + hidden_dim_processor=self.hidden_dim, + hidden_dim_node_encoder=self.hidden_dim, + num_layers_node_encoder=self.num_layers, + hidden_dim_node_decoder=self.hidden_dim, + num_layers_node_decoder=self.num_layers, + hidden_dim_edge_encoder=self.hidden_dim, + num_layers_edge_encoder=self.num_layers, + num_layers_edge_processor=self.num_layers, + num_layers_node_processor=self.num_layers, + aggregation="mean", + num_processor_checkpoint_segments=self.num_processor_checkpoint_segments, + ) + return cast("torch.nn.Module", model) + + def _setup_model_inputs(self, device: "torch.device") -> None: + from torch_geometric.data import Data + + self._device = device + self._shared_edge_index = pnt.mesh_to_edge_index(self._mean_surface) + self._shared_edge_feats = pnt.compute_edge_features( + self._mean_shape_coords, self._shared_edge_index + ) + self._shared_graph = Data( + edge_index=self._shared_edge_index, + num_nodes=len(self._mean_shape_coords), + ) + self._batched_graph_cache = {} + + def _batched_graph(self, batch_len: int) -> tuple[Any, Any]: + """Return (and cache) a batched graph + tiled edge features for ``batch_len``.""" + cached = self._batched_graph_cache.get(batch_len) + if cached is not None: + return cached + from torch_geometric.data import Batch + + assert self._device is not None + graph = Batch.from_data_list([self._shared_graph] * batch_len).to(self._device) + edge_feats = self._shared_edge_feats.repeat(batch_len, 1).to(self._device) + self._batched_graph_cache[batch_len] = (graph, edge_feats) + return graph, edge_feats + + def _forward( + self, model: "torch.nn.Module", node_feats: "torch.Tensor", batch_len: int + ) -> "torch.Tensor": + graph, edge_feats = self._batched_graph(batch_len) + return cast("torch.Tensor", model(node_feats, edge_feats, graph)) + + def _checkpoint_extra(self) -> dict: + return { + "processor_size": self.processor_size, + "hidden_dim": self.hidden_dim, + "num_layers": self.num_layers, + "num_processor_checkpoint_segments": self.num_processor_checkpoint_segments, + "input_dim_edges": 4, + } + + def _save_extra_artifacts(self, output_dir: Path) -> None: + import torch + + torch.save(self._shared_edge_index, output_dir / "shared_edge_index.pt") + torch.save(self._shared_edge_feats, output_dir / "shared_edge_features.pt") + + +class WorkflowTrainPhysicsNeMoMLP(WorkflowTrainPhysicsNeMo): + """Train a PhysicsNeMo :class:`FullyConnected` (MLP) on cardiac mesh stages. + + Each surface point is an independent training row; batches group several + ``(subject, phase)`` samples and shuffle points within the batch to retain + gradient mixing while still streaming from disk. + """ + + _model_tag = "mlp" + _architecture_name = "physicsnemo.models.mlp.FullyConnected" + _shuffle_points_within_batch = True + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.batch_size = 32 # samples per step (points shuffled within the batch) + self.layer_size: int = 512 + self.num_layers: int = 6 + + def set_layer_size(self, layer_size: int) -> None: + """Set the hidden layer width.""" + if layer_size < 1: + raise ValueError(f"layer_size must be >= 1, got {layer_size}") + self.layer_size = layer_size + + def set_num_layers(self, num_layers: int) -> None: + """Set the number of fully connected layers.""" + if num_layers < 1: + raise ValueError(f"num_layers must be >= 1, got {num_layers}") + self.num_layers = num_layers + + def _build_model(self, in_features: int) -> "torch.nn.Module": + + try: + from physicsnemo.models.mlp import FullyConnected + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "The MLP trainer requires PhysicsNeMo, an optional dependency. " + 'Install with: pip install "physiotwin4d[physicsnemo]"' + ) from exc + + model = FullyConnected( + in_features=in_features, + layer_size=self.layer_size, + out_features=3, + num_layers=self.num_layers, + activation_fn="silu", + skip_connections=True, + ) + return cast("torch.nn.Module", model) + + def _setup_model_inputs(self, device: "torch.device") -> None: + # MLP needs no shared graph inputs. + return None + + def _forward( + self, model: "torch.nn.Module", node_feats: "torch.Tensor", batch_len: int + ) -> "torch.Tensor": + return cast("torch.Tensor", model(node_feats)) + + def _checkpoint_extra(self) -> dict: + return {"layer_size": self.layer_size, "num_layers": self.num_layers} + + def _save_extra_artifacts(self, output_dir: Path) -> None: + # No MGN-only graph artifacts for the MLP. + return None diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py index 12fd88c..26d186d 100644 --- a/tests/test_tutorials.py +++ b/tests/test_tutorials.py @@ -77,17 +77,17 @@ def _run_tutorial_script(script_name: str) -> dict[str, Any]: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial01HeartGatedCTToUSD: - """End-to-end test for tutorial_01_heart_gated_ct_to_usd.py.""" +class TestTutorial01aHeartGatedCTToUSD: + """End-to-end test for tutorial_01a_heart_gated_ct_to_usd.py.""" - _class_name = "tutorial_01" + _class_name = "tutorial_01a" def test_run(self, test_directories: dict[str, Path]) -> None: - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_01" - results = _run_tutorial_script("tutorial_01_heart_gated_ct_to_usd.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_01a" + results = _run_tutorial_script("tutorial_01a_heart_gated_ct_to_usd.py") assert results["usd_file"], "USD file path should not be empty" assert Path(results["usd_file"]).exists(), "USD file should exist" - assert results["screenshots"], "Tutorial 1 should produce screenshots" + assert results["screenshots"], "Tutorial 1a should produce screenshots" tt = TestTools( class_name=self._class_name, @@ -123,16 +123,16 @@ def test_run(self, test_directories: dict[str, Path]) -> None: # ----------------------------------------------------------------------------- -# Tutorial 3 - Create Statistical Shape Model +# Tutorial 4a - Create Statistical Shape Model # ----------------------------------------------------------------------------- @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial03CreateStatisticalModel: - """End-to-end test for tutorial_03_create_statistical_model.py.""" +class TestTutorial04aCreateStatisticalModel: + """End-to-end test for tutorial_04a_heart_create_statistical_model.py.""" - _class_name = "tutorial_03_create_statistical_model" + _class_name = "tutorial_04a_heart_create_statistical_model" def test_run(self, test_directories: dict[str, Path]) -> None: kcl_dir = test_directories["data"] / "KCL-Heart-Model" @@ -141,8 +141,8 @@ def test_run(self, test_directories: dict[str, Path]) -> None: "KCL-Heart-Model not downloaded. See data/README.md for instructions." ) - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_03" - results = _run_tutorial_script("tutorial_03_create_statistical_model.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_04a" + results = _run_tutorial_script("tutorial_04a_heart_create_statistical_model.py") assert results["model_file"].exists(), "pca_model.json should exist" assert results["mean_surface_file"].exists(), "Mean surface VTP should exist" @@ -156,10 +156,10 @@ def test_run(self, test_directories: dict[str, Path]) -> None: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial04FitStatisticalModelToPatient: - """End-to-end test for tutorial_04_fit_statistical_model_to_patient.py.""" +class TestTutorial05aFitStatisticalModelToPatient: + """End-to-end test for tutorial_05a_heart_fit_statistical_model_to_patient.py.""" - _class_name = "tutorial_04_fit_statistical_model_to_patient" + _class_name = "tutorial_05a_heart_fit_statistical_model_to_patient" def test_run(self, test_directories: dict[str, Path]) -> None: kcl_dir = test_directories["data"] / "KCL-Heart-Model" @@ -169,18 +169,18 @@ def test_run(self, test_directories: dict[str, Path]) -> None: ) pca_json = ( - _REPO_ROOT / "tutorials" / "output" / "tutorial_03" / "pca_model.json" + _REPO_ROOT / "tutorials" / "output" / "tutorial_04a" / "pca_model.json" ) if not pca_json.exists(): - _run_tutorial_script("tutorial_03_create_statistical_model.py") + _run_tutorial_script("tutorial_04a_heart_create_statistical_model.py") assert pca_json.exists(), ( - "Tutorial 3 bootstrap did not create the expected PCA model file: " + "Tutorial 4a bootstrap did not create the expected PCA model file: " f"{pca_json}" ) - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_04" + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05a" results = _run_tutorial_script( - "tutorial_04_fit_statistical_model_to_patient.py" + "tutorial_05a_heart_fit_statistical_model_to_patient.py" ) assert results["registered_file"].exists(), "Registered VTP should exist" @@ -193,16 +193,16 @@ def test_run(self, test_directories: dict[str, Path]) -> None: # ----------------------------------------------------------------------------- -# Tutorial 5 - VTK to USD +# Tutorial 3 - VTK to USD # ----------------------------------------------------------------------------- @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial05VTKToUSD: - """End-to-end test for tutorial_05_vtk_to_usd.py.""" +class TestTutorial03VTKToUSD: + """End-to-end test for tutorial_03_vtk_to_usd.py.""" - _class_name = "tutorial_05_vtk_to_usd" + _class_name = "tutorial_03_vtk_to_usd" def test_run(self, test_directories: dict[str, Path]) -> None: # Prefer Tutorial 2 output; fall back to any .vtp in data @@ -219,8 +219,8 @@ def test_run(self, test_directories: dict[str, Path]) -> None: ) vtk_file = found[0] - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05" - results = _run_tutorial_script("tutorial_05_vtk_to_usd.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_03" + results = _run_tutorial_script("tutorial_03_vtk_to_usd.py") assert results["usd_file"], "USD file path should not be empty" assert Path(results["usd_file"]).exists(), "USD file should exist" @@ -274,24 +274,24 @@ def _run_eval_tutorial(script_name: str) -> dict[str, Any]: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial08CardiacFitModel: - """End-to-end test for tutorial_08_cardiac_fit_model.py (bring-your-own-data).""" +class TestTutorial08cdBYODFitModel: + """End-to-end test for tutorial_08cd_byod_fit_model_to_patients.py.""" def test_run(self) -> None: if not (_CARDIAC_DATA_ROOT / "duke_data" / "gated_nii").exists(): pytest.skip( - "Cardiac dataset not present at D:/PhysioTwin4D/. Tutorial 8 is " + "Cardiac dataset not present at D:/PhysioTwin4D/. Tutorial 8cd 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" + results = _run_tutorial_script("tutorial_08cd_byod_fit_model_to_patients.py") + assert "patients" in results, "Tutorial 8cd 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.""" +class TestTutorial09cBYODTrainMGN: + """End-to-end test for tutorial_09c_byod_train_physicsnemo_mgn.py.""" def test_run(self) -> None: if not _physicsnemo_available(): @@ -299,30 +299,30 @@ def test_run(self) -> None: 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") + pytest.skip("Tutorial 8cd output not present; run Tutorial 8cd first.") + results = _run_tutorial_script("tutorial_09c_byod_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.""" +class TestTutorial09dBYODTrainMLP: + """End-to-end test for tutorial_09d_byod_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") + pytest.skip("Tutorial 8cd output not present; run Tutorial 8cd first.") + results = _run_tutorial_script("tutorial_09d_byod_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.""" +class TestTutorial10cBYODEvalMGN: + """End-to-end test for tutorial_10c_byod_eval_physicsnemo_mgn.py.""" def test_run(self) -> None: if not _physicsnemo_available(): @@ -332,29 +332,29 @@ def test_run(self) -> None: 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." + "Tutorial 9c checkpoint or cardiac data not present; " + "run Tutorials 8cd and 9c first." ) - results = _run_eval_tutorial("tutorial_10a_cardiac_eval_physicsnemo_mgn.py") - assert "predicted_files" in results + results = _run_eval_tutorial("tutorial_10c_byod_eval_physicsnemo_mgn.py") + assert "predicted_surfaces" in results @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial10bCardiacEvalMLP: - """End-to-end test for tutorial_10b_cardiac_eval_physicsnemo_mlp.py.""" +class TestTutorial10dBYODEvalMLP: + """End-to-end test for tutorial_10d_byod_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" + checkpoint = _TUTORIALS_DIR / "output" / "mlp_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." + "Tutorial 9d checkpoint or cardiac data not present; " + "run Tutorials 8cd and 9d first." ) - results = _run_eval_tutorial("tutorial_10b_cardiac_eval_physicsnemo_mlp.py") - assert "predicted_files" in results + results = _run_eval_tutorial("tutorial_10d_byod_eval_physicsnemo_mlp.py") + assert "predicted_surfaces" in results # ----------------------------------------------------------------------------- diff --git a/tutorials/README.md b/tutorials/README.md index 3e4446a..aaf8939 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -13,24 +13,23 @@ dataset licensing, and expected directory layout. | # | Script | Primary API | Dataset | |---|--------|-------------|---------| -| 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | +| 1a | [tutorial_01a_heart_gated_ct_to_usd.py](tutorial_01a_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | | 2 | [tutorial_02_ct_to_vtk.py](tutorial_02_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | -| 3 | [tutorial_03_create_statistical_model.py](tutorial_03_create_statistical_model.py) | `WorkflowCreateStatisticalModel` | KCL-Heart-Model | -| 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 | +| 3 | [tutorial_03_vtk_to_usd.py](tutorial_03_vtk_to_usd.py) | `WorkflowConvertVTKToUSD` | Output of tutorial 2 | +| 4a | [tutorial_04a_heart_create_statistical_model.py](tutorial_04a_heart_create_statistical_model.py) | `WorkflowCreateStatisticalModel` | KCL-Heart-Model | +| 5a | [tutorial_05a_heart_fit_statistical_model_to_patient.py](tutorial_05a_heart_fit_statistical_model_to_patient.py) | `WorkflowFitStatisticalModelToPatient` | KCL-Heart-Model plus Tutorial 4a output | | 6 | [tutorial_06_reconstruct_highres_4d_ct.py](tutorial_06_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | -| 8 | [tutorial_08_cardiac_fit_model.py](tutorial_08_cardiac_fit_model.py) | `WorkflowFitStatisticalModelToPatient`, `WorkflowReconstructHighres4DCT` | Bring your own (cardiac gated CT, `D:/PhysioTwin4D/`) | -| 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:/PhysioTwin4D/` 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.) +| 8cd | [tutorial_08cd_byod_fit_model_to_patients.py](tutorial_08cd_byod_fit_model_to_patients.py) | `WorkflowFitStatisticalModelToPatient`, `WorkflowReconstructHighres4DCT` | Bring your own (cardiac gated CT, `D:/PhysioTwin4D/`) | +| 9c | [tutorial_09c_byod_train_physicsnemo_mgn.py](tutorial_09c_byod_train_physicsnemo_mgn.py) | `WorkflowTrainPhysicsNeMoMGN` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 8cd output | +| 9d | [tutorial_09d_byod_train_physicsnemo_mlp.py](tutorial_09d_byod_train_physicsnemo_mlp.py) | `WorkflowTrainPhysicsNeMoMLP` (requires `[physicsnemo]` extra) | Tutorial 8cd output | +| 10c | [tutorial_10c_byod_eval_physicsnemo_mgn.py](tutorial_10c_byod_eval_physicsnemo_mgn.py) | `WorkflowInferPhysicsNeMoMGN` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 9c checkpoint | +| 10d | [tutorial_10d_byod_eval_physicsnemo_mlp.py](tutorial_10d_byod_eval_physicsnemo_mlp.py) | `WorkflowInferPhysicsNeMoMLP` (requires `[physicsnemo]` extra) | Tutorial 9d checkpoint | + +> **Tutorials 8cd-10cd are bring-your-own-data.** Unlike the earlier data-driven +> tutorials, they do not use the repository `data/` directory or a downloadable +> sample. Their path constants point at a local `D:/PhysioTwin4D/` cardiac layout +> (gated CT, labelmaps, the KCL volume PCA model, and ICON weights); edit those +> constants to match your own data. ## Running a Tutorial @@ -42,7 +41,7 @@ is read from the repository `data/` directory and outputs are written under ```bash # Run the whole tutorial from the command line -python tutorials/tutorial_01_heart_gated_ct_to_usd.py +python tutorials/tutorial_01a_heart_gated_ct_to_usd.py ``` In VS Code or Cursor, open the tutorial and use **Run Python File** (or run @@ -69,23 +68,23 @@ pytest tests/test_tutorials.py --run-tutorials -v pytest tests/test_tutorials.py --run-tutorials --create-baselines -v # Run a single tutorial test -pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD --run-tutorials -v +pytest tests/test_tutorials.py::TestTutorial01aHeartGatedCTToUSD --run-tutorials -v ``` ## Recommended Order -1. **Tutorial 1** and **Tutorial 2** use Slicer-Heart-CT - prepare it per `data/README.md`, then start here. -2. **Tutorial 5** uses the VTK surfaces produced by Tutorial 2 - run Tutorial 2 first. -3. **Tutorial 3** creates the PCA statistical model from KCL-Heart-Model. -4. **Tutorial 4** applies the statistical model, consuming Tutorial 3 output. +1. **Tutorial 1a** and **Tutorial 2** use Slicer-Heart-CT - prepare it per `data/README.md`, then start here. +2. **Tutorial 3** uses the VTK surfaces produced by Tutorial 2 - run Tutorial 2 first. +3. **Tutorial 4a** creates the PCA statistical model from KCL-Heart-Model. +4. **Tutorial 5a** applies the statistical model, consuming Tutorial 4a output. 5. **Tutorial 6** requires DirLab-4DCT - download it per `data/README.md`. -The cardiac mesh stage-prediction pipeline (Tutorials 8 -> 9 -> 10) is +The cardiac mesh stage-prediction pipeline (Tutorials 8cd -> 9c/9d -> 10c/10d) 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 "physiotwin4d[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. +6. **Tutorial 8cd** 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 9c/9d). +7. **Tutorial 9c / 9d** train a PhysicsNeMo MeshGraphNet (9c) and MLP (9d) to predict a cardiac surface at any cardiac stage. PhysicsNeMo is an optional extra: install with `pip install "physiotwin4d[physicsnemo]"` (requires Python >= 3.11); the MeshGraphNet also needs `torch-geometric`. +8. **Tutorial 10c / 10d** load a trained MeshGraphNet (10c) or MLP (10d) 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_01_heart_gated_ct_to_usd.py b/tutorials/tutorial_01a_heart_gated_ct_to_usd.py similarity index 97% rename from tutorials/tutorial_01_heart_gated_ct_to_usd.py rename to tutorials/tutorial_01a_heart_gated_ct_to_usd.py index 71bd7a9..0bbf857 100644 --- a/tutorials/tutorial_01_heart_gated_ct_to_usd.py +++ b/tutorials/tutorial_01a_heart_gated_ct_to_usd.py @@ -1,5 +1,5 @@ """ -Tutorial 1: Heart-Gated CT to Animated USD +Tutorial 1a: Heart-Gated CT to Animated USD Purpose ------- @@ -94,7 +94,7 @@ DATA_DIR = REPO_ROOT / "data" FULL_DATA_DIR = DATA_DIR / "Slicer-Heart-CT" TEST_DATA_DIR = DATA_DIR / "test" / "slicer_heart_small" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_01" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_01a" LOG_LEVEL = logging.INFO # %% @@ -158,7 +158,7 @@ # %% # Result saving tt = TestTools( - class_name="tutorial_01_heart_gated_ct_to_usd", + class_name="tutorial_01a_heart_gated_ct_to_usd", results_dir=output_dir, log_level=log_level, ) diff --git a/tutorials/tutorial_01b_lung_gated_ct_to_usd.py b/tutorials/tutorial_01b_lung_gated_ct_to_usd.py index 00add8c..8cac3b7 100644 --- a/tutorials/tutorial_01b_lung_gated_ct_to_usd.py +++ b/tutorials/tutorial_01b_lung_gated_ct_to_usd.py @@ -146,6 +146,7 @@ usd_project_name="lung_model", registration_method=registration_method, log_level=log_level, + frames_per_second=4, save_assets=True, ) diff --git a/tutorials/tutorial_02_ct_to_vtk.py b/tutorials/tutorial_02_ct_to_vtk.py index 30369ec..83d9b0b 100644 --- a/tutorials/tutorial_02_ct_to_vtk.py +++ b/tutorials/tutorial_02_ct_to_vtk.py @@ -5,7 +5,7 @@ ------- Segment one 3D CT frame into anatomical groups and save a combined VTK surface file. The output can be inspected directly in PyVista or used as -input for Tutorial 5. +input for Tutorial 3. Data Required ------------- @@ -35,7 +35,7 @@ # 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. +# tutorials also matches the style of tutorial_01a. if __name__ == "__main__": # %% # Data directory specification diff --git a/tutorials/tutorial_05_vtk_to_usd.py b/tutorials/tutorial_03_vtk_to_usd.py similarity index 93% rename from tutorials/tutorial_05_vtk_to_usd.py rename to tutorials/tutorial_03_vtk_to_usd.py index 5b5f94f..15e144e 100644 --- a/tutorials/tutorial_05_vtk_to_usd.py +++ b/tutorials/tutorial_03_vtk_to_usd.py @@ -1,5 +1,5 @@ """ -Tutorial 5: VTK Surface Series to USD +Tutorial 3: VTK Surface Series to USD Purpose ------- @@ -30,7 +30,7 @@ # 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. +# tutorials also matches the style of tutorial_01a. if __name__ == "__main__": # %% # Data directory specification @@ -39,7 +39,7 @@ DATA_DIR = REPO_ROOT / "data" FULL_DATA_DIR = DATA_DIR TEST_DATA_DIR = DATA_DIR / "test" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_05" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_03" BASELINES_DIR = REPO_ROOT / "tests" / "baselines" TUTORIAL_02_SURFACE = ( TUTORIALS_DIR / "output" / "tutorial_02" / "patient_surfaces.vtp" @@ -89,7 +89,7 @@ # %% # Result saving tt = TestTools( - class_name="tutorial_05_vtk_to_usd", + class_name="tutorial_03_vtk_to_usd", results_dir=output_dir, baselines_dir=BASELINES_DIR, log_level=log_level, diff --git a/tutorials/tutorial_03_create_statistical_model.py b/tutorials/tutorial_04a_heart_create_statistical_model.py similarity index 94% rename from tutorials/tutorial_03_create_statistical_model.py rename to tutorials/tutorial_04a_heart_create_statistical_model.py index a411ddf..4ce8989 100644 --- a/tutorials/tutorial_03_create_statistical_model.py +++ b/tutorials/tutorial_04a_heart_create_statistical_model.py @@ -1,10 +1,10 @@ """ -Tutorial 3: Create a PCA Statistical Shape Model +Tutorial 4a: Create a PCA Statistical Shape Model Purpose ------- Build a PCA statistical shape model from a reference mesh and a small population -of sample meshes. Tutorial 4 can reuse the saved ``pca_model.json``. +of sample meshes. Tutorial 5a can reuse the saved ``pca_model.json``. Data Required ------------- @@ -34,7 +34,7 @@ # 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. +# tutorials also matches the style of tutorial_01a. if __name__ == "__main__": # %% # Data directory specification @@ -43,7 +43,7 @@ DATA_DIR = REPO_ROOT / "data" FULL_DATA_DIR = DATA_DIR / "KCL-Heart-Model" TEST_DATA_DIR = DATA_DIR / "test" / "KCL-Heart-Model" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_03" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_04a" BASELINES_DIR = REPO_ROOT / "tests" / "baselines" PCA_COMPONENTS = 10 LOG_LEVEL = logging.INFO @@ -112,7 +112,7 @@ mean_surface.save(str(mean_surface_file)) tt = TestTools( - class_name="tutorial_03_create_statistical_model", + class_name="tutorial_04a_heart_create_statistical_model", results_dir=output_dir, baselines_dir=BASELINES_DIR, log_level=log_level, diff --git a/tutorials/tutorial_04_fit_statistical_model_to_patient.py b/tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py similarity index 91% rename from tutorials/tutorial_04_fit_statistical_model_to_patient.py rename to tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py index 6cb78f2..3088a9d 100644 --- a/tutorials/tutorial_04_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py @@ -1,10 +1,10 @@ """ -Tutorial 4: Fit Statistical Shape Model to Patient Data +Tutorial 5a: Fit Statistical Shape Model to Patient Data Purpose ------- Fit a generic anatomical template mesh to one or more patient-like surface -meshes. If Tutorial 3 has already written ``pca_model.json``, the workflow uses +meshes. If Tutorial 4a has already written ``pca_model.json``, the workflow uses that model to constrain the fitted shape. Data Required @@ -46,15 +46,15 @@ REPO_ROOT = Path(__file__).resolve().parent.parent TUTORIALS_DIR = Path(__file__).resolve().parent DATA_DIR = REPO_ROOT / "data" / "DirLab-4DCT" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_04" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_05a" BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - PCA_JSON = TUTORIALS_DIR / "output" / "tutorial_03" / "pca_model.json" - PCA_MEAN_FILE = TUTORIALS_DIR / "output" / "tutorial_03" / "pca_mean_surface.vtp" + PCA_JSON = TUTORIALS_DIR / "output" / "tutorial_04a" / "pca_model.json" + PCA_MEAN_FILE = TUTORIALS_DIR / "output" / "tutorial_04a" / "pca_mean_surface.vtp" # .mha files are DirLab-4DCT data already converted to HU by # data/DirLab-4DCT/fix_downloaded_data.py. PATIENT_IMAGE_FILE = DATA_DIR / "Case1Pack_T70.mha" SEGMENTATION_METHOD = SegmentChestTotalSegmentator() - SEGMENTATION_METHOD.set_has_highres_heart_license(True) + SEGMENTATION_METHOD.set_has_academic_license(True) # SEGMENTATION_METHOD = SegmentHeartSimplewareTrimmedBranches() # Use when available # and images are contrast-enhanced. # SEGMENTATION_METHOD = SegmentChestTotalSegmentatorWithContrast() # Use when contrast-enhanced @@ -117,6 +117,7 @@ patient_labelmap=heart_labelmap, log_level=log_level, labelmap_interior_object_ids=[141, 142, 143, 144], + # These are the internal chambers of the heart when using TotalSegmentator. ) if pca_model is not None: workflow.set_use_pca_registration( @@ -176,7 +177,7 @@ screenshots.append(after_path) TestTools( - class_name="tutorial_04_fit_statistical_model_to_patient", + class_name="tutorial_05a_heart_fit_statistical_model_to_patient", results_dir=output_dir, baselines_dir=BASELINES_DIR, log_level=log_level, diff --git a/tutorials/tutorial_06_reconstruct_highres_4d_ct.py b/tutorials/tutorial_06_reconstruct_highres_4d_ct.py index 8f33444..ac2c040 100644 --- a/tutorials/tutorial_06_reconstruct_highres_4d_ct.py +++ b/tutorials/tutorial_06_reconstruct_highres_4d_ct.py @@ -34,7 +34,7 @@ # 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. +# tutorials also matches the style of tutorial_01a. if __name__ == "__main__": # %% # Data directory specification @@ -79,6 +79,7 @@ # Workflow initialization registration_method = RegisterImagesGreedyICON(log_level=log_level) registration_method.greedy.set_number_of_iterations(number_of_iterations_Greedy) + registration_method.icon.set_mass_preservation(True) # For non-contrast CT workflow = WorkflowReconstructHighres4DCT( time_series_images=time_series, fixed_image=fixed_image, diff --git a/tutorials/tutorial_08_cardiac_fit_model.py b/tutorials/tutorial_08cd_byod_fit_model_to_patients.py similarity index 95% rename from tutorials/tutorial_08_cardiac_fit_model.py rename to tutorials/tutorial_08cd_byod_fit_model_to_patients.py index c634a00..6c40f4b 100644 --- a/tutorials/tutorial_08_cardiac_fit_model.py +++ b/tutorials/tutorial_08cd_byod_fit_model_to_patients.py @@ -1,13 +1,13 @@ """ -Tutorial 8: Fit the Cardiac SSM and Propagate It Through Gated Phases +Tutorial 8cd: 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: +First stage of the cardiac 4D deep-learning pipeline (Tutorials 08cd -> 09c/09d +-> 10c/10d). For each patient it turns gated CT scans into the +statistical-shape-model (SSM) surfaces and volume meshes that the Tutorial 9c/9d +trainers (``tutorial_09c_byod_train_physicsnemo_mgn.py`` / +``tutorial_09d_byod_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 @@ -25,7 +25,7 @@ Bring Your Own Data ------------------- -This is a bring-your-own-data tutorial. Unlike Tutorials 1-6, it does not use the +This is a bring-your-own-data tutorial. Unlike Tutorials 1-7, it does not use the repository ``data/`` directory or a downloadable sample; the path constants below point at a local ``D:/PhysioTwin4D/`` layout. Edit them to match your own data. @@ -83,7 +83,7 @@ "D:/PhysioTwin4D/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. + # this is also the directory the Tutorial 9c/9d trainers read from. OUTPUT_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") # Simpleware's heart interior chamber labels, excluded from the distance map. LABELMAP_INTERIOR_OBJECT_IDS = [1, 2, 3, 4] @@ -93,7 +93,7 @@ LOG_LEVEL = logging.INFO logging.basicConfig(level=LOG_LEVEL) - logger = logging.getLogger("tutorial_08_cardiac_fit_model") + logger = logging.getLogger("tutorial_08cd_byod_fit_model_to_patients") # In test mode, limit the run to a single patient to keep it tractable. test_mode = TestTools.running_as_test() diff --git a/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py b/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py deleted file mode 100644 index 4cba7a8..0000000 --- a/tutorials/tutorial_09a_cardiac_train_physicsnemo_mgn.py +++ /dev/null @@ -1,878 +0,0 @@ -""" -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:/PhysioTwin4D/`` 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 "physiotwin4d[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 physiotwin4d.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 "physiotwin4d[physicsnemo]"\n' - " pip install torch-geometric" - ) from exc - - -if __name__ == "__main__": - # %% - TUTORIALS_DIR = Path(__file__).resolve().parent - FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") - PCA_MEAN_VTU = Path("D:/PhysioTwin4D/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 deleted file mode 100644 index cc97616..0000000 --- a/tutorials/tutorial_09b_cardiac_train_physicsnemo_mlp.py +++ /dev/null @@ -1,805 +0,0 @@ -""" -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:/PhysioTwin4D/`` 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:/PhysioTwin4D/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 PhysioTwin4D. Install it with:: - - pip install "physiotwin4d[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 physiotwin4d.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 "physiotwin4d[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:/PhysioTwin4D/duke_data/fitted_kcl_meshes") - PCA_MEAN_VTU = Path("D:/PhysioTwin4D/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_09c_byod_train_physicsnemo_mgn.py b/tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py new file mode 100644 index 0000000..ad24c6f --- /dev/null +++ b/tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py @@ -0,0 +1,171 @@ +""" +Tutorial 9c (MGN): Train a PhysicsNeMo MeshGraphNet for cardiac mesh stages. + +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 08cd -> 09c/09d +-> 10c/10d). This tutorial is a thin driver over the reusable +:class:`physiotwin4d.WorkflowTrainPhysicsNeMoMGN` workflow: it discovers the +per-time-point SSM surfaces produced by Tutorial 8cd +(``tutorial_08cd_byod_fit_model_to_patients.py``), writes one JSON manifest per +subject, splits the subjects into train / validation / held-out test, trains the +MeshGraphNet, and evaluates the held-out test subjects with +:class:`physiotwin4d.WorkflowInferPhysicsNeMoMGN`. The companion MLP tutorial is +``tutorial_09d_byod_train_physicsnemo_mlp.py``. + +Why a GNN? +---------- +The SSM mesh has a fixed topology across all subjects and cardiac tissue is a +continuum: adjacent vertices co-vary smoothly. MeshGraphNet encodes that prior +directly by passing messages along mesh edges, giving an explicit +continuum-deformation inductive bias the MLP must infer from coordinates alone. + +Node features (per vertex): [mean_shape_x, mean_shape_y, mean_shape_z, pca_c1 ... pca_cN, stage] +Edge features (per edge): [rel_x, rel_y, rel_z, distance] (from the mean shape) +Output (per vertex): [dx, dy, dz] (displacement in mm) + +Bring Your Own Data +------------------- +The path constants below point at a local ``D:/PhysioTwin4D/`` layout produced by +Tutorial 8cd, not at the repository ``data/`` directory. Edit them to match your +own data location. Run Tutorial 8cd first. + +Extra Install Required +---------------------- +PhysicsNeMo and PyTorch Geometric must be installed:: + + pip install "physiotwin4d[physicsnemo]" +""" + +# %% +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Optional + +from physiotwin4d import WorkflowInferPhysicsNeMoMGN, WorkflowTrainPhysicsNeMoMGN + + +def _gating_stage_from_filename(mesh_file: Path) -> float: + """Extract the normalized cardiac stage [0, 1] from a ``g0TT`` filename stem.""" + 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 _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[Path]: + """Write a per-subject manifest JSON; return its path (or None if incomplete). + + A subject needs a reference SSM surface, a PCA coefficient file, and at least + two gated-phase surfaces. Stages are parsed from the ``g0TT`` phase filenames + and written explicitly into the manifest (the workflow itself never parses + filenames). + """ + sid = subject_dir.name + ref_file = subject_dir / f"{sid}_ssm_surface.vtp" + pca_file = subject_dir / f"{sid}_ssm_pca_coefficients.json" + phase_files = sorted(subject_dir.glob(f"{sid}_g0*_ssm_surface.vtp")) + if not ref_file.exists() or not pca_file.exists() or len(phase_files) < 2: + return None + + manifest = { + "subject_id": sid, + "reference_surface": str(ref_file), + "pca_coefficients": str(pca_file), + "phases": [ + {"surface": str(pf), "stage": _gating_stage_from_filename(pf)} + for pf in phase_files + ], + } + manifests_dir.mkdir(parents=True, exist_ok=True) + manifest_path = manifests_dir / f"{sid}_manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path + + +if __name__ == "__main__": + # %% + TUTORIALS_DIR = Path(__file__).resolve().parent + FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") + PCA_MEAN_VTU = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09c" + MANIFESTS_DIR = TUTORIALS_DIR / "manifests_mgn" + + EPOCHS = 1500 + BATCH_SIZE_GRAPHS = 4 # mini-batch measured in (subject, phase) graphs + LEARNING_RATE = 1.0e-3 + PROCESSOR_SIZE = 3 # message-passing hops + HIDDEN_DIM = 128 + NUM_LAYERS = 2 # MLP layers inside each encoder / processor / decoder block + + # Explicit held-out splits; every other discovered subject is used for training. + TEST_SUBJECTS = ["pm0028"] + VAL_SUBJECTS = ["pm0027"] + LOG_LEVEL = logging.INFO + + def run_tutorial() -> dict[str, Any]: + """Discover subjects, train a MeshGraphNet, and evaluate the test split.""" + logging.basicConfig(level=LOG_LEVEL) + + # Build one manifest per valid subject and partition into splits. + manifests: dict[str, Path] = {} + for subject_dir in sorted(FITTED_MESHES_DIR.glob("pm????")): + manifest_path = _write_subject_manifest(subject_dir, MANIFESTS_DIR) + if manifest_path is not None: + manifests[subject_dir.name] = manifest_path + + if len(manifests) < 3: + raise RuntimeError( + f"Found only {len(manifests)} valid subject(s); need at least 3 " + "for a train / val / test split." + ) + + unknown = [s for s in TEST_SUBJECTS + VAL_SUBJECTS if s not in manifests] + if unknown: + raise ValueError(f"Split subjects not found: {unknown}") + + test_manifests = [manifests[s] for s in TEST_SUBJECTS] + val_manifests = [manifests[s] for s in VAL_SUBJECTS] + train_manifests = [ + p + for sid, p in manifests.items() + if sid not in TEST_SUBJECTS and sid not in VAL_SUBJECTS + ] + logging.info( + "Subject split - train: %d, val: %d, test: %d", + len(train_manifests), + len(val_manifests), + len(test_manifests), + ) + + # Train the MeshGraphNet. + trainer = WorkflowTrainPhysicsNeMoMGN( + train_manifests=train_manifests, + val_manifests=val_manifests, + pca_mean_mesh=PCA_MEAN_VTU, + output_directory=OUTPUT_DIR, + log_level=LOG_LEVEL, + ) + trainer.set_epochs(EPOCHS) + trainer.set_batch_size(BATCH_SIZE_GRAPHS) + trainer.set_learning_rate(LEARNING_RATE) + trainer.set_processor_size(PROCESSOR_SIZE) + trainer.set_hidden_dim(HIDDEN_DIM) + trainer.set_num_layers(NUM_LAYERS) + train_result = trainer.process() + + # Evaluate held-out test subjects against their ground-truth phases. + infer = WorkflowInferPhysicsNeMoMGN( + model_directory=OUTPUT_DIR, log_level=LOG_LEVEL + ) + eval_outputs: dict[str, Any] = {} + for sid in TEST_SUBJECTS: + eval_outputs[sid] = infer.predict( + manifests[sid], output_directory=OUTPUT_DIR / "eval_mgn" / sid + ) + + return {"training": train_result, "evaluation": eval_outputs} + + # %% + tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py b/tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py new file mode 100644 index 0000000..c7dc88c --- /dev/null +++ b/tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py @@ -0,0 +1,167 @@ +""" +Tutorial 9d (MLP): Train a PhysicsNeMo fully connected model for cardiac mesh stages. + +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 08cd -> 09c/09d +-> 10c/10d). This tutorial is a thin driver over the reusable +:class:`physiotwin4d.WorkflowTrainPhysicsNeMoMLP` workflow: it discovers the +per-time-point SSM surfaces produced by Tutorial 8cd +(``tutorial_08cd_byod_fit_model_to_patients.py``), writes one JSON manifest per +subject, splits the subjects into train / validation / held-out test, trains a +FullyConnected (MLP) model, and evaluates the held-out test subjects with +:class:`physiotwin4d.WorkflowInferPhysicsNeMoMLP`. + +The companion Tutorial 9c (``tutorial_09c_byod_train_physicsnemo_mgn.py``) +solves the same task with a MeshGraphNet so the two architectures can be compared +directly; both map a surface point ``(x, y, z, pca_c1 ... pca_cN, stage)`` to its +displacement from the subject's SSM reference surface (the Option B convention), +where the coordinates come from the shared PCA mean shape. + +Batch size note +--------------- +Because the workflow streams samples from disk, ``batch_size`` counts +``(subject, phase)`` samples per step (not individual points as the original +inline trainer did); the MLP shuffles points *within* each batch to keep +gradient mixing. + +Bring Your Own Data +------------------- +The path constants below point at a local ``D:/PhysioTwin4D/`` layout produced by +Tutorial 8cd, not at the repository ``data/`` directory. Edit them to match your +own data location. Run Tutorial 8cd first. + +Extra Install Required +---------------------- +PhysicsNeMo must be installed (requires Python >= 3.11):: + + pip install "physiotwin4d[physicsnemo]" +""" + +# %% +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Optional + +from physiotwin4d import WorkflowInferPhysicsNeMoMLP, WorkflowTrainPhysicsNeMoMLP + + +def _gating_stage_from_filename(mesh_file: Path) -> float: + """Extract the normalized cardiac stage [0, 1] from a ``g0TT`` filename stem.""" + 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 _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[Path]: + """Write a per-subject manifest JSON; return its path (or None if incomplete). + + A subject needs a reference SSM surface, a PCA coefficient file, and at least + two gated-phase surfaces. Stages are parsed from the ``g0TT`` phase filenames + and written explicitly into the manifest (the workflow itself never parses + filenames). + """ + sid = subject_dir.name + ref_file = subject_dir / f"{sid}_ssm_surface.vtp" + pca_file = subject_dir / f"{sid}_ssm_pca_coefficients.json" + phase_files = sorted(subject_dir.glob(f"{sid}_g0*_ssm_surface.vtp")) + if not ref_file.exists() or not pca_file.exists() or len(phase_files) < 2: + return None + + manifest = { + "subject_id": sid, + "reference_surface": str(ref_file), + "pca_coefficients": str(pca_file), + "phases": [ + {"surface": str(pf), "stage": _gating_stage_from_filename(pf)} + for pf in phase_files + ], + } + manifests_dir.mkdir(parents=True, exist_ok=True) + manifest_path = manifests_dir / f"{sid}_manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path + + +if __name__ == "__main__": + # %% + TUTORIALS_DIR = Path(__file__).resolve().parent + FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") + PCA_MEAN_VTU = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09d" + MANIFESTS_DIR = TUTORIALS_DIR / "manifests_mlp" + + EPOCHS = 10000 + BATCH_SIZE_SAMPLES = 32 # (subject, phase) samples per step; points shuffled within + LEARNING_RATE = 1.0e-3 + LAYER_SIZE = 512 + NUM_LAYERS = 6 + + # Explicit held-out splits; every other discovered subject is used for training. + TEST_SUBJECTS = ["pm0028"] + VAL_SUBJECTS = ["pm0027"] + LOG_LEVEL = logging.INFO + + def run_tutorial() -> dict[str, Any]: + """Discover subjects, train an MLP, and evaluate the test split.""" + logging.basicConfig(level=LOG_LEVEL) + + manifests: dict[str, Path] = {} + for subject_dir in sorted(FITTED_MESHES_DIR.glob("pm????")): + manifest_path = _write_subject_manifest(subject_dir, MANIFESTS_DIR) + if manifest_path is not None: + manifests[subject_dir.name] = manifest_path + + if len(manifests) < 3: + raise RuntimeError( + f"Found only {len(manifests)} valid subject(s); need at least 3 " + "for a train / val / test split." + ) + + unknown = [s for s in TEST_SUBJECTS + VAL_SUBJECTS if s not in manifests] + if unknown: + raise ValueError(f"Split subjects not found: {unknown}") + + test_manifests = [manifests[s] for s in TEST_SUBJECTS] + val_manifests = [manifests[s] for s in VAL_SUBJECTS] + train_manifests = [ + p + for sid, p in manifests.items() + if sid not in TEST_SUBJECTS and sid not in VAL_SUBJECTS + ] + logging.info( + "Subject split - train: %d, val: %d, test: %d", + len(train_manifests), + len(val_manifests), + len(test_manifests), + ) + + trainer = WorkflowTrainPhysicsNeMoMLP( + train_manifests=train_manifests, + val_manifests=val_manifests, + pca_mean_mesh=PCA_MEAN_VTU, + output_directory=OUTPUT_DIR, + log_level=LOG_LEVEL, + ) + trainer.set_epochs(EPOCHS) + trainer.set_batch_size(BATCH_SIZE_SAMPLES) + trainer.set_learning_rate(LEARNING_RATE) + trainer.set_layer_size(LAYER_SIZE) + trainer.set_num_layers(NUM_LAYERS) + train_result = trainer.process() + + infer = WorkflowInferPhysicsNeMoMLP( + model_directory=OUTPUT_DIR, log_level=LOG_LEVEL + ) + eval_outputs: dict[str, Any] = {} + for sid in TEST_SUBJECTS: + eval_outputs[sid] = infer.predict( + manifests[sid], output_directory=OUTPUT_DIR / "eval_mlp" / sid + ) + + return {"training": train_result, "evaluation": eval_outputs} + + # %% + tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py b/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py deleted file mode 100644 index 32093ac..0000000 --- a/tutorials/tutorial_10a_cardiac_eval_physicsnemo_mgn.py +++ /dev/null @@ -1,534 +0,0 @@ -""" -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:/PhysioTwin4D/`` 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 "physiotwin4d[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:/PhysioTwin4D/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 deleted file mode 100644 index e006d5d..0000000 --- a/tutorials/tutorial_10b_cardiac_eval_physicsnemo_mlp.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -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:/PhysioTwin4D/`` 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 "physiotwin4d[physicsnemo]"' - ) from exc - -logger = logging.getLogger("tutorial_10b_cardiac_eval_physicsnemo_mlp") - -TUTORIALS_DIR = Path(__file__).resolve().parent -FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/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() diff --git a/tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py b/tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py new file mode 100644 index 0000000..9140ce7 --- /dev/null +++ b/tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py @@ -0,0 +1,135 @@ +""" +Tutorial 10c (MGN): Predict cardiac stage meshes for a subject with a trained +PhysicsNeMo MeshGraphNet. + +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 08cd -> 09c/09d +-> 10c/10d). This tutorial is a thin driver over +:class:`physiotwin4d.WorkflowInferPhysicsNeMoMGN`. It builds a per-subject +manifest from the fitted-meshes directory and calls the workflow to predict +cardiac surfaces for one subject, loading the model trained by Tutorial 9c +(``tutorial_09c_byod_train_physicsnemo_mgn.py``). + +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioTwin4D/`` layout and the Tutorial 9c run directory, not at the +repository ``data/`` directory. + +Usage (command line) +-------------------- + py tutorial_10c_byod_eval_physicsnemo_mgn.py pm0028 --out results/pm0028_mgn + py tutorial_10c_byod_eval_physicsnemo_mgn.py pm0028 --out results/pm0028_mgn --stages 0.0 0.25 0.5 0.75 + +Arguments + subject Subject ID, e.g. pm0028 + --epoch Optional intermittent-checkpoint epoch; omit to use the final + weights stored in the main checkpoint. + --out Output directory (created if missing) + --stages RR-interval fractions to predict (0-1). If omitted, predicts at + every gated phase in the subject's manifest and computes error + statistics against the ground-truth surfaces. When supplied, only + the requested stages are predicted (no error statistics). +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path +from typing import Any, Optional, cast + +from physiotwin4d import WorkflowInferPhysicsNeMoMGN + +logger = logging.getLogger("tutorial_10c_byod_eval_physicsnemo_mgn") + +TUTORIALS_DIR = Path(__file__).resolve().parent +FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") +# Tutorial 9c run directory to evaluate (matches that trainer's OUTPUT_DIR). +MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09c" + +DEFAULT_SUBJECT = "pm0028" +DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10c" / DEFAULT_SUBJECT + + +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 _write_subject_manifest(subject: str, out_dir: Path) -> Path: + """Build a manifest for *subject* from the fitted-meshes directory.""" + 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" + phase_files = sorted(subject_dir.glob(f"{subject}_g0*_ssm_surface.vtp")) + for required in (ref_file, pca_file): + if not required.exists(): + sys.exit(f"Missing: {required}") + if not phase_files: + sys.exit(f"No gated phase files found in {subject_dir}") + + manifest = { + "subject_id": subject, + "reference_surface": str(ref_file), + "pca_coefficients": str(pca_file), + "phases": [ + {"surface": str(pf), "stage": _gating_stage_from_filename(pf)} + for pf in phase_files + ], + } + out_dir.mkdir(parents=True, exist_ok=True) + manifest_path = out_dir / f"{subject}_manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path + + +def predict( + subject: str, + out_dir: Path, + epoch: Optional[int] = None, + stages: Optional[list[float]] = None, +) -> dict[str, Any]: + """Predict cardiac surfaces for *subject* using the trained MeshGraphNet.""" + manifest_path = _write_subject_manifest(subject, out_dir) + infer = WorkflowInferPhysicsNeMoMGN(model_directory=MODEL_DIR, epoch=epoch) + return cast( + "dict[str, Any]", + infer.predict(manifest_path, stages=stages, output_directory=out_dir), + ) + + +def run_tutorial() -> dict[str, Any]: + """Tutorial / test entry point: evaluate DEFAULT_SUBJECT with the final weights.""" + return predict(DEFAULT_SUBJECT, DEFAULT_OUT_DIR) + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Predict cardiac stage meshes for one subject with MeshGraphNet.", + ) + ap.add_argument("subject", help="Subject ID, e.g. pm0028") + ap.add_argument( + "--epoch", type=int, default=None, help="Checkpoint epoch (default: final)" + ) + 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 (omit for per-phase error stats).", + ) + args = ap.parse_args() + predict(args.subject, args.out, epoch=args.epoch, 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_10d_byod_eval_physicsnemo_mlp.py b/tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py new file mode 100644 index 0000000..5e644dc --- /dev/null +++ b/tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py @@ -0,0 +1,135 @@ +""" +Tutorial 10d (MLP): Predict cardiac stage meshes for a subject with a trained +PhysicsNeMo fully connected (MLP) model. + +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 08cd -> 09c/09d +-> 10c/10d). This tutorial is a thin driver over +:class:`physiotwin4d.WorkflowInferPhysicsNeMoMLP`. It builds a per-subject +manifest from the fitted-meshes directory and calls the workflow to predict +cardiac surfaces for one subject, loading the model trained by Tutorial 9d +(``tutorial_09d_byod_train_physicsnemo_mlp.py``). + +This is a bring-your-own-data tutorial: the path constants below point at a local +``D:/PhysioTwin4D/`` layout and the Tutorial 9d run directory, not at the +repository ``data/`` directory. + +Usage (command line) +-------------------- + py tutorial_10d_byod_eval_physicsnemo_mlp.py pm0028 --out results/pm0028_mlp + py tutorial_10d_byod_eval_physicsnemo_mlp.py pm0028 --out results/pm0028_mlp --stages 0.0 0.25 0.5 0.75 + +Arguments + subject Subject ID, e.g. pm0028 + --epoch Optional intermittent-checkpoint epoch; omit to use the final + weights stored in the main checkpoint. + --out Output directory (created if missing) + --stages RR-interval fractions to predict (0-1). If omitted, predicts at + every gated phase in the subject's manifest and computes error + statistics against the ground-truth surfaces. When supplied, only + the requested stages are predicted (no error statistics). +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path +from typing import Any, Optional, cast + +from physiotwin4d import WorkflowInferPhysicsNeMoMLP + +logger = logging.getLogger("tutorial_10d_byod_eval_physicsnemo_mlp") + +TUTORIALS_DIR = Path(__file__).resolve().parent +FITTED_MESHES_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") +# Tutorial 9d run directory to evaluate (matches that trainer's OUTPUT_DIR). +MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09d" + +DEFAULT_SUBJECT = "pm0028" +DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10d" / DEFAULT_SUBJECT + + +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 _write_subject_manifest(subject: str, out_dir: Path) -> Path: + """Build a manifest for *subject* from the fitted-meshes directory.""" + 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" + phase_files = sorted(subject_dir.glob(f"{subject}_g0*_ssm_surface.vtp")) + for required in (ref_file, pca_file): + if not required.exists(): + sys.exit(f"Missing: {required}") + if not phase_files: + sys.exit(f"No gated phase files found in {subject_dir}") + + manifest = { + "subject_id": subject, + "reference_surface": str(ref_file), + "pca_coefficients": str(pca_file), + "phases": [ + {"surface": str(pf), "stage": _gating_stage_from_filename(pf)} + for pf in phase_files + ], + } + out_dir.mkdir(parents=True, exist_ok=True) + manifest_path = out_dir / f"{subject}_manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path + + +def predict( + subject: str, + out_dir: Path, + epoch: Optional[int] = None, + stages: Optional[list[float]] = None, +) -> dict[str, Any]: + """Predict cardiac surfaces for *subject* using the trained MLP.""" + manifest_path = _write_subject_manifest(subject, out_dir) + infer = WorkflowInferPhysicsNeMoMLP(model_directory=MODEL_DIR, epoch=epoch) + return cast( + "dict[str, Any]", + infer.predict(manifest_path, stages=stages, output_directory=out_dir), + ) + + +def run_tutorial() -> dict[str, Any]: + """Tutorial / test entry point: evaluate DEFAULT_SUBJECT with the final weights.""" + return predict(DEFAULT_SUBJECT, DEFAULT_OUT_DIR) + + +def main() -> None: + ap = argparse.ArgumentParser( + description="Predict cardiac stage meshes for one subject with an MLP.", + ) + ap.add_argument("subject", help="Subject ID, e.g. pm0028") + ap.add_argument( + "--epoch", type=int, default=None, help="Checkpoint epoch (default: final)" + ) + 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 (omit for per-phase error stats).", + ) + args = ap.parse_args() + predict(args.subject, args.out, epoch=args.epoch, 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()