diff --git a/docs/architecture.rst b/docs/architecture.rst index 951a644..8f2a6e8 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -91,29 +91,29 @@ Primary Workflows AI Surrogate Workflows (PhysicsNeMo) ===================================== -The final tier of tutorials (``tutorials/tutorial_08cd`` through -``tutorial_10d``) turns a fitted statistical shape model into a trained AI +The final tier of tutorials (``tutorials/tutorial_08`` through +``tutorial_10``) 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_08cd_byod_fit_model_to_patients.py`` +``tutorial_08_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_09c_byod_train_physicsnemo_mgn.py`` / -``tutorial_09d_byod_train_physicsnemo_mlp.py`` +``tutorial_09_byod_train_physicsnemo_mgn.py`` / +``tutorial_09_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 + ``WorkflowTrainPhysicsNeMoMLP`` — 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_10c_byod_eval_physicsnemo_mgn.py`` / -``tutorial_10d_byod_eval_physicsnemo_mlp.py`` +``tutorial_10_byod_eval_physicsnemo_mgn.py`` / +``tutorial_10_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 diff --git a/docs/index.rst b/docs/index.rst index caa1e97..4921168 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -22,8 +22,8 @@
- - 1a + + 01

Heart-Gated CT to Animated USD

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

Slicer-Heart-CT @@ -40,17 +40,17 @@

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

Tutorial 2 output - - 4a + + 04

Create a PCA Shape Model

Build a statistical shape model from aligned cardiac meshes.

KCL-Heart-Model - - 5a + + 05

Fit Statistical Model to Patient

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

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

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

DirLab-4DCT - - 8cd + + 08

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 - - 9cd + + 09

Train a PhysicsNeMo Cardiac Stage Model

-

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

- Tutorial 8cd output +

Train a PhysicsNeMo MeshGraphNet or MLP to predict cardiac meshes at requested stages.

+ Tutorial 8 output - - 10cd + + 10

Predict and Evaluate Cardiac Surfaces

-

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

- Tutorial 9c / 9d output +

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

+ Tutorial 9 output
diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 17e2d8b..ca23189 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_01a_...py``) or cell-by-cell in VS Code or Cursor. +(``python tutorials/tutorial_01_...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,9 +41,9 @@ tutorials: .. code-block:: bash - python tutorials/tutorial_01a_heart_gated_ct_to_usd.py + python tutorials/tutorial_01_heart_gated_ct_to_usd.py - python tutorials/tutorial_02_ct_to_vtk.py + python tutorials/tutorial_02_heart_ct_to_vtk.py Tutorial paths are defined near the top of each script. To use different paths, edit the script constants or use the installed ``physiotwin4d-*`` CLI commands. @@ -52,16 +52,16 @@ recommended run order. Recommended run order: -1. Tutorials 1a and 2 first, after downloading Slicer-Heart-CT data. +1. Tutorials 1 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. +3. Tutorial 4 after downloading KCL-Heart-Model. +4. Tutorial 5 after Tutorial 4 because it can consume Tutorial 4 output. 5. Tutorial 6 after downloading DirLab-4DCT (manual). -6. Tutorial 8cd after preparing your own cardiac gated CT, labelmaps, KCL volume +6. Tutorial 8 after preparing your own cardiac gated CT, labelmaps, KCL volume PCA model, and ICON weights (bring-your-own-data). -7. Tutorial 9c and/or 9d after Tutorial 8cd because they train from its fitted +7. Tutorial 9 (MGN and/or MLP) after Tutorial 8 because they train from its fitted meshes. -8. Tutorial 10c and/or 10d after Tutorial 9c / 9d because they evaluate the +8. Tutorial 10 (MGN and/or MLP) after Tutorial 9 because they evaluate the trained checkpoints. Prerequisites diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 61e17eb..9c5e594 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -19,8 +19,8 @@ Tutorials
- - 1a + + 01

Heart-Gated CT to Animated USD

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

Slicer-Heart-CT @@ -37,17 +37,17 @@ Tutorials

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

Tutorial 2 output - - 4a + + 04

Create a PCA Shape Model

Build a statistical shape model from aligned cardiac meshes.

KCL-Heart-Model - - 5a + + 05

Fit Statistical Model to Patient

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

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

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

DirLab-4DCT - - 8cd + + 08

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 - - 9cd + + 09

Train a PhysicsNeMo Cardiac Stage Model

-

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

- Tutorial 8cd output +

Train a PhysicsNeMo MeshGraphNet or MLP to predict cardiac meshes at requested stages.

+ Tutorial 8 output - - 10cd + + 10

Predict and Evaluate Cardiac Surfaces

-

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

- Tutorial 9c / 9d output +

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

+ Tutorial 9 output
@@ -93,8 +93,8 @@ 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 8cd-10cd are -bring-your-own-data (see the note before Tutorial 8cd) and do not use a +``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 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 1a and 2 after downloading Slicer-Heart-CT data (above). +1. Run Tutorials 1 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. +3. Run Tutorial 4 after downloading KCL-Heart-Model (above). +4. Run Tutorial 5 after Tutorial 4 because it can consume the PCA model output. 5. Run Tutorial 6 after downloading DirLab-4DCT (manual — see above). -6. Run Tutorial 8cd after preparing your own cardiac gated CT, labelmaps, KCL +6. Run Tutorial 8 after preparing your own cardiac gated CT, labelmaps, KCL volume PCA model, and ICON weights (bring-your-own-data; see the note below). -7. Run Tutorial 9c and/or 9d after Tutorial 8cd because they train from its +7. Run Tutorial 9 (MGN and/or MLP) after Tutorial 8 because they train from its fitted meshes. -8. Run Tutorial 10c and/or 10d after Tutorial 9c / 9d because they evaluate +8. Run Tutorial 10 (MGN and/or MLP) after Tutorial 9 because they evaluate the trained checkpoints. -Tutorial 1a: Heart-Gated CT to Animated USD -=========================================== +Tutorial 1: Heart-Gated CT to Animated USD +========================================== Script - ``tutorials/tutorial_01a_heart_gated_ct_to_usd.py`` + ``tutorials/tutorial_01_heart_gated_ct_to_usd.py`` Workflow ``WorkflowConvertImageToUSD`` @@ -168,7 +168,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_01a_heart_gated_ct_to_usd.py + python tutorials/tutorial_01_heart_gated_ct_to_usd.py Outputs Registered phase images, transformed contours, preview screenshots, and an @@ -178,7 +178,7 @@ Tutorial 2: CT Segmentation to VTK Surfaces =========================================== Script - ``tutorials/tutorial_02_ct_to_vtk.py`` + ``tutorials/tutorial_02_heart_ct_to_vtk.py`` Workflow ``WorkflowConvertImageToVTK`` @@ -227,16 +227,16 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_02_ct_to_vtk.py + python tutorials/tutorial_02_heart_ct_to_vtk.py Outputs Segmentation artifacts, VTK PolyData surfaces, and preview screenshots. -Tutorial 4a: Create a PCA Shape Model -===================================== +Tutorial 4: Create a PCA Shape Model +==================================== Script - ``tutorials/tutorial_04a_heart_create_statistical_model.py`` + ``tutorials/tutorial_04_heart_create_statistical_model.py`` Workflow ``WorkflowCreateStatisticalModel`` @@ -273,16 +273,16 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_04a_heart_create_statistical_model.py + python tutorials/tutorial_04_heart_create_statistical_model.py Outputs PCA model files, mean shape, and component diagnostics. -Tutorial 5a: Fit Statistical Model to Patient -============================================= +Tutorial 5: Fit Statistical Model to Patient +============================================ Script - ``tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py`` + ``tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py`` Workflow ``WorkflowFitStatisticalModelToPatient`` @@ -321,7 +321,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py + python tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py Outputs Patient-fitted statistical model surfaces and registration diagnostics. @@ -330,7 +330,7 @@ Tutorial 3: VTK Surface Series to Animated USD ============================================== Script - ``tutorials/tutorial_03_vtk_to_usd.py`` + ``tutorials/tutorial_03_heart_vtk_to_usd.py`` Workflow ``WorkflowConvertVTKToUSD`` @@ -378,7 +378,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_03_vtk_to_usd.py + python tutorials/tutorial_03_heart_vtk_to_usd.py Outputs Time-sampled USD scene and conversion logs for Omniverse inspection. @@ -387,7 +387,7 @@ Tutorial 6: Reconstruct High-Resolution 4D CT ============================================= Script - ``tutorials/tutorial_06_reconstruct_highres_4d_ct.py`` + ``tutorials/tutorial_06_lung_reconstruct_highres_4d_ct.py`` Workflow ``WorkflowReconstructHighres4DCT`` @@ -431,7 +431,7 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_06_reconstruct_highres_4d_ct.py + python tutorials/tutorial_06_lung_reconstruct_highres_4d_ct.py Outputs Registered respiratory phases, reconstructed high-resolution CT volumes, @@ -439,18 +439,18 @@ Outputs .. note:: - Tutorials 8cd-10cd form the cardiac mesh stage-prediction pipeline and are + Tutorials 8-10 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 8cd: Fit the Cardiac SSM and Propagate Through Gated Phases -=================================================================== +Tutorial 8: Fit the Cardiac SSM and Propagate Through Gated Phases +================================================================== Script - ``tutorials/tutorial_08cd_byod_fit_model_to_patients.py`` + ``tutorials/tutorial_08_byod_fit_model_to_patients.py`` Workflow ``WorkflowFitStatisticalModelToPatient`` (PCA registration) and @@ -507,24 +507,24 @@ Inner API usage Run .. code-block:: bash - python tutorials/tutorial_08cd_byod_fit_model_to_patients.py + python tutorials/tutorial_08_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 9c / 9d: Train a PhysicsNeMo Cardiac Stage Model -========================================================= +Tutorial 9: Train a PhysicsNeMo Cardiac Stage Model +=================================================== Script - ``tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py`` (MeshGraphNet) and - ``tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py`` (MLP) + ``tutorials/tutorial_09_byod_train_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_09_byod_train_physicsnemo_mlp.py`` (MLP) Inner API usage These are thin drivers over the reusable - ``WorkflowTrainPhysicsNeMoMGN`` (9c) and ``WorkflowTrainPhysicsNeMoMLP`` (9d) + ``WorkflowTrainPhysicsNeMoMGN`` (MGN) and ``WorkflowTrainPhysicsNeMoMLP`` (MLP) workflows, which train ``physicsnemo.models.meshgraphnet.MeshGraphNet`` or - ``physicsnemo.models.mlp.FullyConnected`` on Tutorial 8cd's fitted meshes. The + ``physicsnemo.models.mlp.FullyConnected`` on Tutorial 8'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 @@ -532,7 +532,7 @@ Inner API usage the rest of the workflow layer. Dataset - Tutorial 8cd fitted-mesh outputs. + Tutorial 8 fitted-mesh outputs. Preview .. figure:: assets/example.gif @@ -555,29 +555,29 @@ Extra install Run .. code-block:: bash - python tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py - python tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py + python tutorials/tutorial_09_byod_train_physicsnemo_mgn.py + python tutorials/tutorial_09_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 10c / 10d: Predict and Evaluate Cardiac Surfaces -========================================================= +Tutorial 10: Predict and Evaluate Cardiac Surfaces +================================================== Script - ``tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py`` (MeshGraphNet) and - ``tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py`` (MLP) + ``tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py`` (MeshGraphNet) and + ``tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py`` (MLP) Inner API usage Thin drivers over ``WorkflowInferPhysicsNeMoMGN`` / ``WorkflowInferPhysicsNeMoMLP`` - that load a Tutorial 9c/9d checkpoint and predict cardiac surfaces for one + that load a Tutorial 9 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 9c / 9d trained checkpoints plus the Tutorial 8cd fitted meshes. + Tutorial 9 trained checkpoints plus the Tutorial 8 fitted meshes. Preview .. figure:: assets/example.gif @@ -595,7 +595,7 @@ Preview Run .. code-block:: bash - python tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py pm0002 --out results/pm0002 + python tutorials/tutorial_10_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_OUT_DIR`` constants. diff --git a/pyproject.toml b/pyproject.toml index 9df685b..958a800 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -112,12 +112,16 @@ cuda13 = [ # PhysicsNeMo pulls in a large dependency stack (CUDA-only toolchain, narrower # 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. +# graph batching / connectivity utilities used by the MeshGraphNet (MGN) path, +# and ``torch-scatter`` provides its scatter/segment reduction backend. Install +# with ``pip install physiotwin4d[physicsnemo]`` when you need it. Note: +# nvidia-physicsnemo itself requires Python >= 3.11. On a custom NVIDIA PyTorch +# build with no matching pre-built wheel, torch-scatter must be built from +# source: ``pip install torch-scatter --no-build-isolation``. physicsnemo = [ "nvidia-physicsnemo>=2.0.0", "torch-geometric>=2.5.0", + "torch-scatter>=2.1.0", ] dev = [ "bumpver>=2023.0.0", @@ -274,11 +278,11 @@ module = [ "physiotwin4d.cli.visualize_pca_modes", "physiotwin4d.vtk_to_usd.mesh_utils", "physiotwin4d.vtk_to_usd.vtk_reader", - "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", + "tutorial_08_byod_fit_model_to_patients", + "tutorial_09_byod_train_physicsnemo_mgn", + "tutorial_09_byod_train_physicsnemo_mlp", + "tutorial_10_byod_eval_physicsnemo_mgn", + "tutorial_10_byod_eval_physicsnemo_mlp", ] disable_error_code = ["import-not-found", "import-untyped"] diff --git a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py index 0d09819..fee1144 100644 --- a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py @@ -96,9 +96,12 @@ def main() -> int: ) parser.add_argument( "--register-reference", - action="store_true", - default=False, - help="Register reference frame to fixed image (default: use identity)", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Register the reference time frame to the reference image " + "(default: enabled; use --no-register-reference for an identity transform)" + ), ) parser.add_argument( "--prior-weight", @@ -285,9 +288,9 @@ def main() -> int: workflow = WorkflowReconstructHighres4DCT( time_series_images=time_series_images, - fixed_image=fixed_image, - reference_frame=args.reference_frame, - register_reference=args.register_reference, + reference_image=fixed_image, + reference_time_frame=args.reference_frame, + register_reference_time_frame_to_reference_image=args.register_reference, registration_method=registration_method, ) diff --git a/src/physiotwin4d/segment_chest_total_segmentator.py b/src/physiotwin4d/segment_chest_total_segmentator.py index f2480c3..098d79d 100644 --- a/src/physiotwin4d/segment_chest_total_segmentator.py +++ b/src/physiotwin4d/segment_chest_total_segmentator.py @@ -269,7 +269,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: background regions as the skin outline. Note: - Requires GPU acceleration (device="gpu") for reasonable performance. + Requires GPU acceleration (device="gpu:0") for reasonable performance. The method automatically handles coordinate system conversions between ITK and nibabel formats. @@ -293,11 +293,11 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: # nr_thr_resamp defaults to 1; TotalSegmentator's post-prediction # resampling back to native resolution is CPU-bound and benefits # from parallelizing across the available cores. - resamp_threads = min(8, os.cpu_count() or 1) + resamp_threads = min(12, os.cpu_count() or 1) output_nib_image_total = totalsegmentator( nib_image, task="total", - device="gpu", + device="gpu:0", fast=self.fast_mode, nr_thr_resamp=resamp_threads, ) @@ -311,7 +311,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: output_nib_image_heart = totalsegmentator( nib_image, task="heartchambers_highres", - device="gpu", + device="gpu:0", nr_thr_resamp=resamp_threads, ) labelmap_arr_heart = output_nib_image_heart.get_fdata().astype( @@ -333,7 +333,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: output_nib_image_tissue_4_types = totalsegmentator( nib_image, task="tissue_4_types", - device="gpu", + device="gpu:0", nr_thr_resamp=resamp_threads, ) labelmap_arr_tissue_4_types = ( @@ -360,7 +360,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: output_nib_image_lung = totalsegmentator( nib_image, task="lung_vessels", - device="gpu", + device="gpu:0", nr_thr_resamp=resamp_threads, ) labelmap_arr_lung = output_nib_image_lung.get_fdata().astype(np.uint8) @@ -374,7 +374,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: self.log_info("Running body task") output_nib_image_body = totalsegmentator( - nib_image, task="body", device="gpu", nr_thr_resamp=resamp_threads + nib_image, task="body", device="gpu:0", nr_thr_resamp=resamp_threads ) labelmap_arr_body = output_nib_image_body.get_fdata().astype(np.uint8) # labelmap_arr_body contains: 1=body, 2=body_trunc, 3=body_extremities, diff --git a/src/physiotwin4d/workflow_convert_image_to_usd.py b/src/physiotwin4d/workflow_convert_image_to_usd.py index 5d3dd6c..1e2bae7 100644 --- a/src/physiotwin4d/workflow_convert_image_to_usd.py +++ b/src/physiotwin4d/workflow_convert_image_to_usd.py @@ -179,7 +179,7 @@ def process(self) -> dict[str, str]: return painted_usd_files - def _register_with_mask( + def _register( self, fixed_image: itk.Image, fixed_mask: Optional[itk.Image], @@ -268,17 +268,20 @@ def _segment_and_register_frames(self) -> None: moving_image = self.time_series_images[i] - moving_segmentation_results = self.segmenter.segment(moving_image) - moving_labelmap = moving_segmentation_results["labelmap"] - if self.save_assets: - itk.imwrite( - moving_labelmap, - os.path.join(self.output_directory, f"slice_{i:03d}_labelmap.mha"), - compression=True, - ) - if len(self.dynamic_labelmap_ids) > 0: self.registrar.set_fixed_mask(reference_dynamic_mask) + + moving_segmentation_results = self.segmenter.segment(moving_image) + moving_labelmap = moving_segmentation_results["labelmap"] + if self.save_assets: + itk.imwrite( + moving_labelmap, + os.path.join( + self.output_directory, f"slice_{i:03d}_labelmap.mha" + ), + compression=True, + ) + moving_labelmap_arr = itk.GetArrayFromImage(moving_labelmap) moving_is_dynamic_arr = np.isin( moving_labelmap_arr, self.dynamic_labelmap_ids @@ -298,7 +301,7 @@ def _segment_and_register_frames(self) -> None: compression=True, ) - dynamic_reg_results = self._register_with_mask( + dynamic_reg_results = self._register( self.reference_image, reference_dynamic_mask, moving_image, @@ -315,7 +318,7 @@ def _segment_and_register_frames(self) -> None: static_labelmap, self.mask_dilation_radius ) - static_reg_results = self._register_with_mask( + static_reg_results = self._register( self.reference_image, reference_static_mask, moving_image, @@ -330,7 +333,7 @@ def _segment_and_register_frames(self) -> None: } ) else: - all_reg_results = self._register_with_mask( + all_reg_results = self._register( self.reference_image, None, moving_image, diff --git a/src/physiotwin4d/workflow_infer_physicsnemo.py b/src/physiotwin4d/workflow_infer_physicsnemo.py index 15310b0..019311b 100644 --- a/src/physiotwin4d/workflow_infer_physicsnemo.py +++ b/src/physiotwin4d/workflow_infer_physicsnemo.py @@ -146,9 +146,10 @@ def _load_weights(self, epoch: Optional[int]) -> dict: ) 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) - ) + ckpt = torch.load(str(epoch_file), map_location="cpu", weights_only=True) + # Self-describing checkpoints wrap the weights under "model_state_dict"; + # bare/legacy epoch checkpoints are the state dict itself. + return cast(dict, ckpt.get("model_state_dict", ckpt)) def _predict_displacements( self, pca_coeffs: np.ndarray, stage: float diff --git a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py index 49d9ab9..a6c9695 100644 --- a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py @@ -49,7 +49,7 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): **Input Requirements:** - time_series_images: Ordered list of 3D images (typically lower resolution) - - fixed_image: High-resolution reference image + - reference_image: High-resolution reference image - All images should be in the same anatomical coordinate system ``registration_method`` accepts a pre-configured @@ -60,16 +60,17 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): Attributes: time_series_images (list[itk.Image]): Ordered list of time-series images - fixed_image (itk.Image): High-resolution reference image - reference_frame (int): Index of reference frame in time series - register_reference (bool): Whether to register reference frame + reference_image (itk.Image): High-resolution reference image + reference_time_frame (int): Index of reference time frame in time series + register_reference_time_frame_to_reference_image (bool): Whether to register + the reference time frame to the reference image prior_weight (float): Weight for temporal smoothing (0.0-1.0) upsample_to_fixed_resolution (bool): Whether to upsample reconstruction registrar (RegisterTimeSeriesImages): Internal registration object forward_transforms (list[itk.Transform]): one per frame; each warps its - moving image onto the fixed grid + moving image onto the reference grid inverse_transforms (list[itk.Transform]): one per frame; each warps the - fixed image onto that frame's moving grid (used for reconstruction) + reference image onto that frame's moving grid (used for reconstruction) losses (list[float]): Registration loss values reconstructed_images (list[itk.Image]): Reconstructed high-resolution images @@ -77,8 +78,8 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): >>> # Initialize workflow with data >>> workflow = WorkflowReconstructHighres4DCT( ... time_series_images=lowres_images, - ... fixed_image=highres_reference, - ... reference_frame=7, + ... reference_image=highres_reference, + ... reference_time_frame=7, ... ) >>> >>> # Access results @@ -90,9 +91,9 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): def __init__( self, time_series_images: list[itk.Image], - fixed_image: itk.Image, - reference_frame: int = 0, - register_reference: bool = False, + reference_image: itk.Image, + reference_time_frame: int = 0, + register_reference_time_frame_to_reference_image: bool = True, registration_method: Optional[RegisterImagesBase] = None, log_level: int | str = logging.INFO, ): @@ -101,13 +102,13 @@ def __init__( Args: time_series_images (list[itk.Image]): Ordered list of 3D time-series images to be registered and reconstructed - fixed_image (itk.Image): High-resolution 3D reference image - reference_frame (int, optional): Index of the reference frame in the - time series. Registration proceeds bidirectionally from this frame. + reference_image (itk.Image): High-resolution 3D reference image + reference_time_frame (int, optional): Index of the reference time frame in + the time series. Registration proceeds bidirectionally from this frame. Default: 0 - register_reference (bool, optional): If True, register the reference frame - to the fixed image. If False, use identity transform for reference. - Default: False + register_reference_time_frame_to_reference_image (bool, optional): If True, + register the reference time frame to the reference image. If False, use + an identity transform for that frame. Default: True registration_method (Optional[RegisterImagesBase]): Registration backend instance. Defaults to a new :class:`RegisterImagesGreedyICON` when None. @@ -116,7 +117,7 @@ def __init__( Raises: ValueError: If time_series_images is empty - ValueError: If reference_frame is out of range + ValueError: If reference_time_frame is out of range TypeError: If registration_method is neither None nor a RegisterImagesBase instance """ @@ -129,9 +130,9 @@ def __init__( if not time_series_images: raise ValueError("time_series_images cannot be empty") - if reference_frame < 0 or reference_frame >= len(time_series_images): + if reference_time_frame < 0 or reference_time_frame >= len(time_series_images): raise ValueError( - f"reference_frame {reference_frame} out of range " + f"reference_time_frame {reference_time_frame} out of range " f"[0, {len(time_series_images) - 1}]" ) @@ -144,9 +145,11 @@ def __init__( # Store input data self.time_series_images = time_series_images - self.fixed_image = fixed_image - self.reference_frame = reference_frame - self.register_reference = register_reference + self.reference_image = reference_image + self.reference_time_frame = reference_time_frame + self.register_reference_time_frame_to_reference_image = ( + register_reference_time_frame_to_reference_image + ) # Initialize parameters with defaults self.prior_weight: float = 0.0 @@ -238,9 +241,9 @@ def register_time_series(self) -> dict: Returns: dict: Dictionary containing: - 'forward_transforms' (list[itk.Transform]): one per frame; - each warps its moving image onto the fixed grid + each warps its moving image onto the reference grid - 'inverse_transforms' (list[itk.Transform]): one per frame; - each warps the fixed image onto that frame's moving grid + each warps the reference image onto that frame's moving grid (see docs/developer/transform_conventions) - 'losses' (list[float]): Registration loss value for each image @@ -252,23 +255,26 @@ def register_time_series(self) -> dict: ) # Configure registrar - self.registrar.set_fixed_image(self.fixed_image) + self.registrar.set_fixed_image(self.reference_image) self.registrar.set_modality(self.modality) self.registrar.set_mask_dilation(self.mask_dilation_mm) self.registrar.set_fixed_mask(self.fixed_mask) self.log_info(f"Registration method: {type(self.registrar.registrar).__name__}") self.log_info(f"Number of time points: {len(self.time_series_images)}") - self.log_info(f"Reference frame: {self.reference_frame}") - self.log_info(f"Register reference: {self.register_reference}") + self.log_info(f"Reference time frame: {self.reference_time_frame}") + self.log_info( + "Register reference time frame to reference image: " + f"{self.register_reference_time_frame_to_reference_image}" + ) self.log_info(f"Prior weight: {self.prior_weight}") # Perform registration result = self.registrar.register_time_series( moving_images=self.time_series_images, moving_masks=self.moving_masks, - reference_frame=self.reference_frame, - register_reference=self.register_reference, + reference_frame=self.reference_time_frame, + register_reference=self.register_reference_time_frame_to_reference_image, prior_weight=self.prior_weight, ) @@ -371,7 +377,7 @@ def process(self) -> dict: self.log_info(f" Number of time points: {len(self.time_series_images)}") registrar_type = type(self.registrar.registrar).__name__ self.log_info(f" Registration method: {registrar_type}") - self.log_info(f" Reference frame: {self.reference_frame}") + self.log_info(f" Reference time frame: {self.reference_time_frame}") self.log_info(f" Prior weight: {self.prior_weight}") self.log_info( f" Upsample to fixed resolution: {self.upsample_to_fixed_resolution}" diff --git a/src/physiotwin4d/workflow_train_physicsnemo.py b/src/physiotwin4d/workflow_train_physicsnemo.py index 430f5ac..9f5fa79 100644 --- a/src/physiotwin4d/workflow_train_physicsnemo.py +++ b/src/physiotwin4d/workflow_train_physicsnemo.py @@ -312,7 +312,11 @@ 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: + # Inherit the exact stats when the checkpoint carries them (final models + # and, since this change, periodic epoch checkpoints). Bare/legacy epoch + # checkpoints hold only weights: recompute from the data, which is + # identical for an unchanged subject set (the normal resume case). + if resume_ckpt is not None and "coordinate_mean" in resume_ckpt: return { "coordinate_mean": np.array(resume_ckpt["coordinate_mean"], np.float32), "coordinate_scale": np.array( @@ -322,6 +326,11 @@ def _compute_normalization( "pca_scale": np.array(resume_ckpt["pca_scale"], np.float32), "displacement_scale": float(resume_ckpt["displacement_scale"]), } + if resume_ckpt is not None: + self.log_warning( + "Resume checkpoint has no normalization stats (bare weights-only " + "checkpoint); recomputing them from the current data." + ) coord = self._mean_shape_coords coordinate_mean = coord.mean(axis=0) @@ -504,7 +513,7 @@ def _train( output_dir / f"{self._model_tag}_stage_model_epoch_{epoch + 1:05d}.pt" ) - torch.save(pnt.uncompiled_state_dict(model), ckpt_path) + torch.save(self._build_checkpoint(model, stats), ckpt_path) self.log_info( " intermittent test epoch %05d/%d train RMSE=%.4f mm " "val RMSE=%.4f mm checkpoint=%s", @@ -553,6 +562,29 @@ def _evaluate_rmse( model.train() return float(np.sqrt(total_sq / max(n_points, 1))) + def _build_checkpoint( + self, model: "torch.nn.Module", stats: dict + ) -> dict[str, Any]: + """Assemble a self-describing checkpoint (weights + normalization stats). + + Both the periodic epoch checkpoints and the final model share this + payload so training can resume from — and inference can load — any saved + checkpoint, not just the final one. + """ + checkpoint: dict[str, Any] = { + "model_state_dict": pnt.uncompiled_state_dict(model), + "architecture": self._architecture_name, + "in_features": 3 + int(stats["pca_mean"].shape[0]) + 1, + "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"], + } + checkpoint.update(self._checkpoint_extra()) + return checkpoint + def _save_model( self, model: "torch.nn.Module", @@ -574,21 +606,10 @@ def _save_model( 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()) + checkpoint = self._build_checkpoint(model, stats) + checkpoint["train_subject_ids"] = train_ids + checkpoint["val_subject_ids"] = val_ids + checkpoint["resumed_from"] = str(self.resume_from) if self.resume_from else None torch.save(checkpoint, checkpoint_file) n_pca = int(stats["pca_mean"].shape[0]) diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py index 26d186d..f565b37 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 TestTutorial01aHeartGatedCTToUSD: - """End-to-end test for tutorial_01a_heart_gated_ct_to_usd.py.""" +class TestTutorial01HeartGatedCTToUSD: + """End-to-end test for tutorial_01_heart_gated_ct_to_usd.py.""" - _class_name = "tutorial_01a" + _class_name = "tutorial_01_heart_gated_ct_to_usd" def test_run(self, test_directories: dict[str, Path]) -> None: - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_01a" - results = _run_tutorial_script("tutorial_01a_heart_gated_ct_to_usd.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_01_heart" + results = _run_tutorial_script("tutorial_01_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 1a should produce screenshots" + assert results["screenshots"], "Tutorial 1 should produce screenshots" tt = TestTools( class_name=self._class_name, @@ -104,14 +104,14 @@ def test_run(self, test_directories: dict[str, Path]) -> None: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial02CTToVTK: - """End-to-end test for tutorial_02_ct_to_vtk.py.""" +class TestTutorial02HeartCTToVTK: + """End-to-end test for tutorial_02_heart_ct_to_vtk.py.""" - _class_name = "tutorial_02_ct_to_vtk" + _class_name = "tutorial_02_heart_ct_to_vtk" def test_run(self, test_directories: dict[str, Path]) -> None: - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_02" - results = _run_tutorial_script("tutorial_02_ct_to_vtk.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_02_heart" + results = _run_tutorial_script("tutorial_02_heart_ct_to_vtk.py") assert results["surface_file"].exists(), "Combined VTP surface should exist" tt = TestTools( @@ -123,16 +123,16 @@ def test_run(self, test_directories: dict[str, Path]) -> None: # ----------------------------------------------------------------------------- -# Tutorial 4a - Create Statistical Shape Model +# Tutorial 4 - Create Statistical Shape Model # ----------------------------------------------------------------------------- @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial04aCreateStatisticalModel: - """End-to-end test for tutorial_04a_heart_create_statistical_model.py.""" +class TestTutorial04CreateStatisticalModel: + """End-to-end test for tutorial_04_heart_create_statistical_model.py.""" - _class_name = "tutorial_04a_heart_create_statistical_model" + _class_name = "tutorial_04_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_04a" - results = _run_tutorial_script("tutorial_04a_heart_create_statistical_model.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_04_heart" + results = _run_tutorial_script("tutorial_04_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 TestTutorial05aFitStatisticalModelToPatient: - """End-to-end test for tutorial_05a_heart_fit_statistical_model_to_patient.py.""" +class TestTutorial05FitStatisticalModelToPatient: + """End-to-end test for tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py.""" - _class_name = "tutorial_05a_heart_fit_statistical_model_to_patient" + _class_name = "tutorial_05_heart_to_lung_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_04a" / "pca_model.json" + _REPO_ROOT / "tutorials" / "output" / "tutorial_04_heart" / "pca_model.json" ) if not pca_json.exists(): - _run_tutorial_script("tutorial_04a_heart_create_statistical_model.py") + _run_tutorial_script("tutorial_04_heart_create_statistical_model.py") assert pca_json.exists(), ( - "Tutorial 4a bootstrap did not create the expected PCA model file: " + "Tutorial 4 bootstrap did not create the expected PCA model file: " f"{pca_json}" ) - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05a" + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05_heart_to_lung" results = _run_tutorial_script( - "tutorial_05a_heart_fit_statistical_model_to_patient.py" + "tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py" ) assert results["registered_file"].exists(), "Registered VTP should exist" @@ -199,15 +199,19 @@ def test_run(self, test_directories: dict[str, Path]) -> None: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial03VTKToUSD: - """End-to-end test for tutorial_03_vtk_to_usd.py.""" +class TestTutorial03HeartVTKToUSD: + """End-to-end test for tutorial_03_heart_vtk_to_usd.py.""" - _class_name = "tutorial_03_vtk_to_usd" + _class_name = "tutorial_03_heart_vtk_to_usd" def test_run(self, test_directories: dict[str, Path]) -> None: # Prefer Tutorial 2 output; fall back to any .vtp in data tutorial2_vtp = ( - _REPO_ROOT / "tutorials" / "output" / "tutorial_02" / "patient_surfaces.vtp" + _REPO_ROOT + / "tutorials" + / "output" + / "tutorial_02_heart" + / "patient_surfaces.vtp" ) vtk_file = tutorial2_vtp if tutorial2_vtp.exists() else None if vtk_file is None: @@ -219,8 +223,8 @@ def test_run(self, test_directories: dict[str, Path]) -> None: ) vtk_file = found[0] - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_03" - results = _run_tutorial_script("tutorial_03_vtk_to_usd.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_03_heart" + results = _run_tutorial_script("tutorial_03_heart_vtk_to_usd.py") assert results["usd_file"], "USD file path should not be empty" assert Path(results["usd_file"]).exists(), "USD file should exist" @@ -274,24 +278,24 @@ def _run_eval_tutorial(script_name: str) -> dict[str, Any]: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial08cdBYODFitModel: - """End-to-end test for tutorial_08cd_byod_fit_model_to_patients.py.""" +class TestTutorial08BYODFitModel: + """End-to-end test for tutorial_08_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 8cd is " + "Cardiac dataset not present at D:/PhysioTwin4D/. Tutorial 8 is " "bring-your-own-data; see tutorials/README.md." ) - results = _run_tutorial_script("tutorial_08cd_byod_fit_model_to_patients.py") - assert "patients" in results, "Tutorial 8cd should report processed patients" + results = _run_tutorial_script("tutorial_08_byod_fit_model_to_patients.py") + assert "patients" in results, "Tutorial 8 should report processed patients" @pytest.mark.tutorial @pytest.mark.slow @pytest.mark.requires_gpu -class TestTutorial09cBYODTrainMGN: - """End-to-end test for tutorial_09c_byod_train_physicsnemo_mgn.py.""" +class TestTutorial09BYODTrainMGN: + """End-to-end test for tutorial_09_byod_train_physicsnemo_mgn.py.""" def test_run(self) -> None: if not _physicsnemo_available(): @@ -299,61 +303,65 @@ 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 8cd output not present; run Tutorial 8cd first.") - results = _run_tutorial_script("tutorial_09c_byod_train_physicsnemo_mgn.py") + pytest.skip("Tutorial 8 output not present; run Tutorial 8 first.") + results = _run_tutorial_script("tutorial_09_byod_train_physicsnemo_mgn.py") assert isinstance(results, dict) @pytest.mark.tutorial @pytest.mark.slow @pytest.mark.requires_gpu -class TestTutorial09dBYODTrainMLP: - """End-to-end test for tutorial_09d_byod_train_physicsnemo_mlp.py.""" +class TestTutorial09BYODTrainMLP: + """End-to-end test for tutorial_09_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 8cd output not present; run Tutorial 8cd first.") - results = _run_tutorial_script("tutorial_09d_byod_train_physicsnemo_mlp.py") + pytest.skip("Tutorial 8 output not present; run Tutorial 8 first.") + results = _run_tutorial_script("tutorial_09_byod_train_physicsnemo_mlp.py") assert isinstance(results, dict) @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial10cBYODEvalMGN: - """End-to-end test for tutorial_10c_byod_eval_physicsnemo_mgn.py.""" +class TestTutorial10BYODEvalMGN: + """End-to-end test for tutorial_10_byod_eval_physicsnemo_mgn.py.""" def test_run(self) -> None: if not _physicsnemo_available(): pytest.skip("PhysicsNeMo not installed (optional [physicsnemo] extra).") if not _torch_geometric_available(): pytest.skip("torch-geometric not installed (required for MeshGraphNet).") - checkpoint = _TUTORIALS_DIR / "output_mgn" / "mgn_stage_model.pt" + checkpoint = ( + _TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" / "mgn_stage_model.pt" + ) if not checkpoint.exists() or not _CARDIAC_FITTED_MESHES_DIR.exists(): pytest.skip( - "Tutorial 9c checkpoint or cardiac data not present; " - "run Tutorials 8cd and 9c first." + "Tutorial 9 checkpoint or cardiac data not present; " + "run Tutorials 8 and 9 first." ) - results = _run_eval_tutorial("tutorial_10c_byod_eval_physicsnemo_mgn.py") + results = _run_eval_tutorial("tutorial_10_byod_eval_physicsnemo_mgn.py") assert "predicted_surfaces" in results @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial10dBYODEvalMLP: - """End-to-end test for tutorial_10d_byod_eval_physicsnemo_mlp.py.""" +class TestTutorial10BYODEvalMLP: + """End-to-end test for tutorial_10_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" / "mlp_stage_model.pt" + checkpoint = ( + _TUTORIALS_DIR / "output" / "tutorial_09_byod_mlp" / "mlp_stage_model.pt" + ) if not checkpoint.exists() or not _CARDIAC_FITTED_MESHES_DIR.exists(): pytest.skip( - "Tutorial 9d checkpoint or cardiac data not present; " - "run Tutorials 8cd and 9d first." + "Tutorial 9 checkpoint or cardiac data not present; " + "run Tutorials 8 and 9 first." ) - results = _run_eval_tutorial("tutorial_10d_byod_eval_physicsnemo_mlp.py") + results = _run_eval_tutorial("tutorial_10_byod_eval_physicsnemo_mlp.py") assert "predicted_surfaces" in results @@ -364,10 +372,10 @@ def test_run(self) -> None: @pytest.mark.tutorial @pytest.mark.slow -class TestTutorial06ReconstructHighres4DCT: - """End-to-end test for tutorial_06_reconstruct_highres_4d_ct.py.""" +class TestTutorial06LungReconstructHighres4DCT: + """End-to-end test for tutorial_06_lung_reconstruct_highres_4d_ct.py.""" - _class_name = "tutorial_06_reconstruct_highres_4d_ct" + _class_name = "tutorial_06_lung_reconstruct_highres_4d_ct" def test_run(self, test_directories: dict[str, Path]) -> None: dirlab_dir = test_directories["data"] / "DirLab-4DCT" / "Case1" @@ -376,8 +384,8 @@ def test_run(self, test_directories: dict[str, Path]) -> None: "DirLab-4DCT Case1 not downloaded. See data/README.md for instructions." ) - out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_06" - results = _run_tutorial_script("tutorial_06_reconstruct_highres_4d_ct.py") + out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_06_lung" + results = _run_tutorial_script("tutorial_06_lung_reconstruct_highres_4d_ct.py") assert results["reconstructed_files"], ( "At least one reconstructed frame expected" ) diff --git a/tests/test_workflow_reconstruct_highres_4d_ct.py b/tests/test_workflow_reconstruct_highres_4d_ct.py index 1eb36b6..1c15c74 100644 --- a/tests/test_workflow_reconstruct_highres_4d_ct.py +++ b/tests/test_workflow_reconstruct_highres_4d_ct.py @@ -25,7 +25,7 @@ def test_default_registration_method_is_greedy_icon() -> None: matching this workflow's historical 'Greedy_ICON' string default.""" workflow = WorkflowReconstructHighres4DCT( time_series_images=[_small_image(), _small_image()], - fixed_image=_small_image(), + reference_image=_small_image(), ) assert isinstance(workflow.registrar.registrar, RegisterImagesGreedyICON) @@ -35,7 +35,7 @@ def test_registration_method_rejects_wrong_type() -> None: with pytest.raises(TypeError, match="registration_method must be"): WorkflowReconstructHighres4DCT( time_series_images=[_small_image(), _small_image()], - fixed_image=_small_image(), + reference_image=_small_image(), registration_method="ICON", # type: ignore[arg-type] ) @@ -45,7 +45,7 @@ def test_caller_supplied_instance_is_used_as_is() -> None: registrar: RegisterImagesBase = RegisterImagesICON() workflow = WorkflowReconstructHighres4DCT( time_series_images=[_small_image(), _small_image()], - fixed_image=_small_image(), + reference_image=_small_image(), registration_method=registrar, ) assert workflow.registrar.registrar is registrar diff --git a/tutorials/README.md b/tutorials/README.md index aaf8939..ad58866 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -13,19 +13,20 @@ dataset licensing, and expected directory layout. | # | Script | Primary API | Dataset | |---|--------|-------------|---------| -| 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_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) | -| 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 +| 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | +| 1 | [tutorial_01_lung_gated_ct_to_usd.py](tutorial_01_lung_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Lung gated 4D CT (prepare first) | +| 2 | [tutorial_02_heart_ct_to_vtk.py](tutorial_02_heart_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | +| 3 | [tutorial_03_heart_vtk_to_usd.py](tutorial_03_heart_vtk_to_usd.py) | `WorkflowConvertVTKToUSD` | Output of tutorial 2 | +| 4 | [tutorial_04_heart_create_statistical_model.py](tutorial_04_heart_create_statistical_model.py) | `WorkflowCreateStatisticalModel` | KCL-Heart-Model | +| 5 | [tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py](tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py) | `WorkflowFitStatisticalModelToPatient` | KCL-Heart-Model plus Tutorial 4 output | +| 6 | [tutorial_06_lung_reconstruct_highres_4d_ct.py](tutorial_06_lung_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | +| 8 | [tutorial_08_byod_fit_model_to_patients.py](tutorial_08_byod_fit_model_to_patients.py) | `WorkflowFitStatisticalModelToPatient`, `WorkflowReconstructHighres4DCT` | Bring your own (cardiac gated CT, `D:/PhysioTwin4D/`) | +| 9 | [tutorial_09_byod_train_physicsnemo_mgn.py](tutorial_09_byod_train_physicsnemo_mgn.py) | `WorkflowTrainPhysicsNeMoMGN` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 8 output | +| 9 | [tutorial_09_byod_train_physicsnemo_mlp.py](tutorial_09_byod_train_physicsnemo_mlp.py) | `WorkflowTrainPhysicsNeMoMLP` (requires `[physicsnemo]` extra) | Tutorial 8 output | +| 10 | [tutorial_10_byod_eval_physicsnemo_mgn.py](tutorial_10_byod_eval_physicsnemo_mgn.py) | `WorkflowInferPhysicsNeMoMGN` (requires `[physicsnemo]` extra + `torch-geometric`) | Tutorial 9 checkpoint | +| 10 | [tutorial_10_byod_eval_physicsnemo_mlp.py](tutorial_10_byod_eval_physicsnemo_mlp.py) | `WorkflowInferPhysicsNeMoMLP` (requires `[physicsnemo]` extra) | Tutorial 9 checkpoint | + +> **Tutorials 8-10 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 @@ -41,7 +42,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_01a_heart_gated_ct_to_usd.py +python tutorials/tutorial_01_heart_gated_ct_to_usd.py ``` In VS Code or Cursor, open the tutorial and use **Run Python File** (or run @@ -68,23 +69,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::TestTutorial01aHeartGatedCTToUSD --run-tutorials -v +pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD --run-tutorials -v ``` ## Recommended Order -1. **Tutorial 1a** and **Tutorial 2** use Slicer-Heart-CT - prepare it per `data/README.md`, then start here. +1. **Tutorial 1** 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. +3. **Tutorial 4** creates the PCA statistical model from KCL-Heart-Model. +4. **Tutorial 5** applies the statistical model, consuming Tutorial 4 output. 5. **Tutorial 6** requires DirLab-4DCT - download it per `data/README.md`. -The cardiac mesh stage-prediction pipeline (Tutorials 8cd -> 9c/9d -> 10c/10d) is +The cardiac mesh stage-prediction pipeline (Tutorials 8 -> 9 -> 10) is bring-your-own-data and runs in order: -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. +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 9 (MGN / MLP)** train a PhysicsNeMo MeshGraphNet and MLP 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 10 (MGN / MLP)** load a trained MeshGraphNet or MLP 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_01a_heart_gated_ct_to_usd.py b/tutorials/tutorial_01_heart_gated_ct_to_usd.py similarity index 85% rename from tutorials/tutorial_01a_heart_gated_ct_to_usd.py rename to tutorials/tutorial_01_heart_gated_ct_to_usd.py index 0bbf857..36fe603 100644 --- a/tutorials/tutorial_01a_heart_gated_ct_to_usd.py +++ b/tutorials/tutorial_01_heart_gated_ct_to_usd.py @@ -1,5 +1,5 @@ """ -Tutorial 1a: Heart-Gated CT to Animated USD +Tutorial 1: Heart-Gated CT to Animated USD Purpose ------- @@ -66,7 +66,6 @@ download notebook or download the file manually before running this tutorial. """ -# %% # Imports from __future__ import annotations @@ -75,47 +74,50 @@ import itk -from physiotwin4d.register_images_icon import RegisterImagesICON -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_convert_image_to_usd import ( +from physiotwin4d import ( + RegisterImagesICON, + SegmentChestTotalSegmentatorWithContrast, + TestTools, WorkflowConvertImageToUSD, ) +# Only run if this script is not imported as a module + # nnUNetv2 (used by TotalSegmentator inside WorkflowConvertImageToUSD) # spawns a multiprocessing.Pool. On Windows the spawn start method re-imports # this script in each child; without the __name__ == "__main__" guard around # the top-level work, that re-import fires workflow.process() again and # Python's spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - 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_01a" - LOG_LEVEL = logging.INFO - - # %% - # Data reading - test_mode = TestTools.running_as_test() + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - log_level = LOG_LEVEL + class_name = "tutorial_01_heart_gated_ct_to_usd" - output_dir.mkdir(parents=True, exist_ok=True) + output_dir = tutorials_dir / "output" / "tutorial_01_heart" + test_mode = TestTools.running_as_test() if test_mode: + data_dir = repo_root / "data" / "test" / "slicer_heart_small" number_of_registration_iterations = 1 + frame_files = sorted(data_dir.glob("slice_???.mha"))[0:2] else: + data_dir = repo_root / "data" / "Slicer-Heart-CT" number_of_registration_iterations = 10 + frame_files = sorted(data_dir.glob("slice_???.mha")) - # %% - frame_files = sorted(data_dir.glob("slice_???.mha")) - if test_mode: - frame_files = frame_files[:2] + log_level = logging.INFO + + registration_method = RegisterImagesICON(log_level=log_level) + registration_method.set_number_of_iterations(number_of_registration_iterations) + + segmentation_method = SegmentChestTotalSegmentatorWithContrast(log_level=log_level) + segmentation_method.set_has_academic_license(True) + + # Directory setup and data reading + + output_dir.mkdir(parents=True, exist_ok=True) input_filenames = [str(path) for path in frame_files] if not input_filenames: @@ -131,10 +133,7 @@ print("Number of time-series images:", len(time_series_images)) - # %% # Workflow initialization - registration_method = RegisterImagesICON(log_level=log_level) - registration_method.set_number_of_iterations(number_of_registration_iterations) workflow = WorkflowConvertImageToUSD( time_series_images=time_series_images, @@ -142,23 +141,23 @@ output_directory=str(output_dir), usd_project_name="cardiac_model", registration_method=registration_method, + segmentation_method=segmentation_method, log_level=log_level, save_assets=True, ) - # %% # Workflow execution - usd_files = workflow.process() + workflow_results = workflow.process() + # if dynamic_labelmap_ids is not None, there are two USD files if len(workflow.dynamic_labelmap_ids) > 0: - usd_file = output_dir / usd_files["dynamic"] + usd_file = output_dir / workflow_results["dynamic"] else: - usd_file = output_dir / usd_files["all"] + usd_file = output_dir / workflow_results["all"] - # %% # Result saving tt = TestTools( - class_name="tutorial_01a_heart_gated_ct_to_usd", + class_name=class_name, results_dir=output_dir, log_level=log_level, ) diff --git a/tutorials/tutorial_01b_lung_gated_ct_to_usd.py b/tutorials/tutorial_01_lung_gated_ct_to_usd.py similarity index 78% rename from tutorials/tutorial_01b_lung_gated_ct_to_usd.py rename to tutorials/tutorial_01_lung_gated_ct_to_usd.py index 8cac3b7..7f60bc6 100644 --- a/tutorials/tutorial_01b_lung_gated_ct_to_usd.py +++ b/tutorials/tutorial_01_lung_gated_ct_to_usd.py @@ -1,5 +1,5 @@ """ -Tutorial 1b: Lung-Gated 4D CT to Animated USD +Tutorial 1: Lung-Gated 4D CT to Animated USD Purpose ------- @@ -67,7 +67,6 @@ ``data/DirLab-4DCT/fix_downloaded_data.py`` before running this tutorial. """ -# %% # Imports from __future__ import annotations @@ -76,49 +75,52 @@ import itk -from physiotwin4d.register_images_icon import RegisterImagesICON -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_convert_image_to_usd import ( +from physiotwin4d import ( + RegisterImagesICON, + SegmentChestTotalSegmentator, + TestTools, WorkflowConvertImageToUSD, ) +# Only run if this script is not imported as a module + # nnUNetv2 (used by TotalSegmentator inside WorkflowConvertImageToUSD) # spawns a multiprocessing.Pool. On Windows the spawn start method re-imports # this script in each child; without the __name__ == "__main__" guard around # the top-level work, that re-import fires workflow.process() again and # Python's spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - DATA_DIR = REPO_ROOT / "data" - FULL_DATA_DIR = DATA_DIR / "DirLab-4DCT" - TEST_DATA_DIR = DATA_DIR / "test" / "DirLab-4DCT" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_01b" - LOG_LEVEL = logging.INFO - - # %% - # Data reading - test_mode = TestTools.running_as_test() + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - log_level = LOG_LEVEL + class_name = "tutorial_01_lung_gated_ct_to_usd" - output_dir.mkdir(parents=True, exist_ok=True) + output_dir = tutorials_dir / "output" / "tutorial_01_lung" + # .mha files are DirLab-4DCT data already converted to HU by + # data/DirLab-4DCT/fix_downloaded_data.py. + test_mode = TestTools.running_as_test() if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" number_of_registration_iterations = 1 + frame_files = sorted(data_dir.glob("Case1Pack_T??.mha"))[0:2] else: + data_dir = repo_root / "data" / "DirLab-4DCT" number_of_registration_iterations = 10 + frame_files = sorted(data_dir.glob("Case1Pack_T??.mha")) - # %% - # .mha files are DirLab-4DCT data already converted to HU by - # data/DirLab-4DCT/fix_downloaded_data.py. - frame_files = sorted(data_dir.glob("Case1Pack_T??.mha")) - if test_mode: - frame_files = frame_files[:2] + log_level = logging.INFO + + registration_method = RegisterImagesICON(log_level=log_level) + registration_method.set_number_of_iterations(number_of_registration_iterations) + + segmentation_method = SegmentChestTotalSegmentator(log_level=log_level) + segmentation_method.set_has_academic_license(True) + + # Directory setup and data reading + + output_dir.mkdir(parents=True, exist_ok=True) input_filenames = [str(path) for path in frame_files] if not input_filenames: @@ -134,10 +136,7 @@ print("Number of time-series images:", len(time_series_images)) - # %% # Workflow initialization - registration_method = RegisterImagesICON(log_level=log_level) - registration_method.set_number_of_iterations(number_of_registration_iterations) workflow = WorkflowConvertImageToUSD( time_series_images=time_series_images, @@ -145,24 +144,24 @@ output_directory=str(output_dir), usd_project_name="lung_model", registration_method=registration_method, + segmentation_method=segmentation_method, log_level=log_level, frames_per_second=4, save_assets=True, ) - # %% # Workflow execution - usd_files = workflow.process() + workflow_results = workflow.process() + # if dynamic_labelmap_ids is not None, there are two USD files if len(workflow.dynamic_labelmap_ids) > 0: - usd_file = output_dir / usd_files["dynamic"] + usd_file = output_dir / workflow_results["dynamic"] else: - usd_file = output_dir / usd_files["all"] + usd_file = output_dir / workflow_results["all"] - # %% # Result saving tt = TestTools( - class_name="tutorial_01b_lung_gated_ct_to_usd", + class_name=class_name, results_dir=output_dir, log_level=log_level, ) @@ -170,7 +169,7 @@ screenshots: list[Path] = [] test_image_num = int(0.7 * len(input_filenames)) - test_image_path = output_dir / f"slice_{test_image_num:03d}_registered.mha" + test_image_path = output_dir / f"slice_{test_image_num:03d}_all_registered.mha" if test_image_path.exists(): test_image = itk.imread(str(test_image_path)) screenshots.append( @@ -185,22 +184,6 @@ ) ) - test_labelmap_path = output_dir / f"slice_{test_image_num:03d}_labelmap.mha" - if test_labelmap_path.exists(): - test_labelmap = itk.imread(str(test_labelmap_path)) - screenshots.append( - tt.save_screenshot_image_slice( - test_image, - f"slice_{test_image_num:03d}_labelmap_test.png", - axis=0, - slice_fraction=0.5, - colormap="gray", - vmin=-200, - vmax=600, - overlay_mask=test_labelmap, - ) - ) - if usd_file.exists(): screenshots.append( tt.save_screenshot_openusd( diff --git a/tutorials/tutorial_02_ct_to_vtk.py b/tutorials/tutorial_02_heart_ct_to_vtk.py similarity index 72% rename from tutorials/tutorial_02_ct_to_vtk.py rename to tutorials/tutorial_02_heart_ct_to_vtk.py index 83d9b0b..11c7119 100644 --- a/tutorials/tutorial_02_ct_to_vtk.py +++ b/tutorials/tutorial_02_heart_ct_to_vtk.py @@ -13,7 +13,6 @@ Test data: ``data/test/slicer_heart_small/slice_???.mha`` """ -# %% # Imports from __future__ import annotations @@ -23,48 +22,52 @@ import itk import pyvista as pv -from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.segment_chest_total_segmentator_with_contrast import ( +from physiotwin4d import ( + ContourTools, SegmentChestTotalSegmentatorWithContrast, + TestTools, + WorkflowConvertImageToVTK, ) -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK + +# Only run if this script is not imported as a module # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this # script in each child; without the __name__ == "__main__" guard around # top-level work, that re-import fires the segmenter again and Python's -# spawn-cascade detector raises RuntimeError. Wrapping consistently across -# tutorials also matches the style of tutorial_01a. +# spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - 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_02" - BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - LOG_LEVEL = logging.INFO + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_02_heart_ct_to_vtk" + + output_dir = tutorials_dir / "output" / "tutorial_02_heart" # In addition to the combined surface file always saved below, also # save one VTP per anatomy group (e.g. heart.vtp, lung.vtp) and/or one # VTP per individual anatomical structure (e.g. left_ventricle.vtp). - SAVE_GROUP_SURFACES = True - SAVE_LABEL_SURFACES = True + save_group_surfaces = True + save_label_surfaces = True - # %% - # Data reading test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "slicer_heart_small" + else: + data_dir = repo_root / "data" / "Slicer-Heart-CT" + + frame_files = sorted(data_dir.glob("slice_???.mha")) - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - log_level = LOG_LEVEL + log_level = logging.INFO + + segmentation_method = SegmentChestTotalSegmentatorWithContrast(log_level=log_level) + segmentation_method.set_has_academic_license(True) + + # Directory setup and data reading output_dir.mkdir(parents=True, exist_ok=True) - frame_files = sorted(data_dir.glob("slice_???.mha")) if not frame_files: raise FileNotFoundError( "Slicer-Heart-CT frame data not found. Checked:\n" @@ -75,26 +78,22 @@ ct_file = frame_files[0] ct_image = itk.imread(str(ct_file)) - # %% # Workflow initialization + workflow = WorkflowConvertImageToVTK( - segmentation_method=SegmentChestTotalSegmentatorWithContrast( - log_level=log_level - ), + segmentation_method=segmentation_method, log_level=log_level, ) - # %% # Workflow execution # # surface_target_reduction decimates each exported VTP surface. result = workflow.process( input_image=ct_image, surface_target_reduction=0.5, - extract_label_surfaces=SAVE_LABEL_SURFACES, + extract_label_surfaces=save_label_surfaces, ) - # %% # Result saving surface_file = Path( ContourTools.save_combined_surface( @@ -103,21 +102,21 @@ prefix="patient", ) ) - if SAVE_GROUP_SURFACES: + if save_group_surfaces: ContourTools.save_surfaces( result["surfaces"], str(output_dir), prefix="patient" ) - if SAVE_LABEL_SURFACES: + if save_label_surfaces: ContourTools.save_surfaces( result["label_surfaces"], str(output_dir), prefix="patient" ) labelmap_file = output_dir / "patient_labelmap.mha" itk.imwrite(result["labelmap"], str(labelmap_file), compression=True) + # Testing tt = TestTools( - class_name="tutorial_02_ct_to_vtk", + class_name=class_name, results_dir=output_dir, - baselines_dir=BASELINES_DIR, log_level=log_level, ) diff --git a/tutorials/tutorial_03_vtk_to_usd.py b/tutorials/tutorial_03_heart_vtk_to_usd.py similarity index 61% rename from tutorials/tutorial_03_vtk_to_usd.py rename to tutorials/tutorial_03_heart_vtk_to_usd.py index 15e144e..09063a1 100644 --- a/tutorials/tutorial_03_vtk_to_usd.py +++ b/tutorials/tutorial_03_heart_vtk_to_usd.py @@ -8,11 +8,10 @@ Data Required ------------- -Preferred input: ``tutorials/output/tutorial_02/patient_surfaces.vtp`` +Preferred input: ``tutorials/output/tutorial_02_heart/patient_surfaces.vtp`` Fallback input: any ``*.vtp`` under ``data`` or ``data/test`` """ -# %% # Imports from __future__ import annotations @@ -22,44 +21,49 @@ import pyvista as pv -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_convert_vtk_to_usd import WorkflowConvertVTKToUSD +from physiotwin4d import ( + TestTools, + WorkflowConvertVTKToUSD, +) + +# Only run if this script is not imported as a module # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this # script in each child; without the __name__ == "__main__" guard around # top-level work, that re-import fires the segmenter again and Python's -# spawn-cascade detector raises RuntimeError. Wrapping consistently across -# tutorials also matches the style of tutorial_01a. +# spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - DATA_DIR = REPO_ROOT / "data" - FULL_DATA_DIR = DATA_DIR - TEST_DATA_DIR = DATA_DIR / "test" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_03" - BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - TUTORIAL_02_SURFACE = ( - TUTORIALS_DIR / "output" / "tutorial_02" / "patient_surfaces.vtp" + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_03_heart_vtk_to_usd" + + output_dir = tutorials_dir / "output" / "tutorial_03_heart" + baselines_dir = repo_root / "tests" / "baselines" + + # Preferred input: the combined surface saved by Tutorial 2. Leave vtk_file as + # None to auto-discover (Tutorial 2 output first, then any *.vtp under data_dir). + tutorial_02_surface = ( + tutorials_dir / "output" / "tutorial_02_heart" / "patient_surfaces.vtp" ) - VTK_FILE: Optional[Path] = None - LOG_LEVEL = logging.INFO + vtk_file: Optional[Path] = None - # %% - # Data reading test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" + else: + data_dir = repo_root / "data" + + log_level = logging.INFO - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - vtk_file = VTK_FILE - log_level = LOG_LEVEL + # Directory setup and data reading output_dir.mkdir(parents=True, exist_ok=True) - if vtk_file is None and TUTORIAL_02_SURFACE.exists(): - vtk_file = TUTORIAL_02_SURFACE + if vtk_file is None and tutorial_02_surface.exists(): + vtk_file = tutorial_02_surface if vtk_file is None: vtk_candidates = sorted(data_dir.rglob("*.vtp")) if not vtk_candidates: @@ -69,9 +73,10 @@ ) vtk_file = vtk_candidates[0] - # %% - # Workflow initialization mesh = pv.read(str(vtk_file)) + + # Workflow initialization + workflow = WorkflowConvertVTKToUSD( input_meshes=[mesh], usd_project_name="surfaces", @@ -82,16 +87,14 @@ log_level=log_level, ) - # %% # Workflow execution usd_file = workflow.process() - # %% - # Result saving + # Testing tt = TestTools( - class_name="tutorial_03_vtk_to_usd", + class_name=class_name, results_dir=output_dir, - baselines_dir=BASELINES_DIR, + baselines_dir=baselines_dir, log_level=log_level, ) diff --git a/tutorials/tutorial_04a_heart_create_statistical_model.py b/tutorials/tutorial_04_heart_create_statistical_model.py similarity index 58% rename from tutorials/tutorial_04a_heart_create_statistical_model.py rename to tutorials/tutorial_04_heart_create_statistical_model.py index 4ce8989..46242b6 100644 --- a/tutorials/tutorial_04a_heart_create_statistical_model.py +++ b/tutorials/tutorial_04_heart_create_statistical_model.py @@ -1,10 +1,10 @@ """ -Tutorial 4a: Create a PCA Statistical Shape Model +Tutorial 4: 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 5a can reuse the saved ``pca_model.json``. +of sample meshes. Tutorial 5 can reuse the saved ``pca_model.json``. Data Required ------------- @@ -12,7 +12,6 @@ Test data: ``data/test/KCL-Heart-Model`` """ -# %% # Imports from __future__ import annotations @@ -24,42 +23,39 @@ import numpy as np import pyvista as pv -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_create_statistical_model import ( +from physiotwin4d import ( + TestTools, WorkflowCreateStatisticalModel, ) +# Only run if this script is not imported as a module + # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this # script in each child; without the __name__ == "__main__" guard around # top-level work, that re-import fires the segmenter again and Python's -# spawn-cascade detector raises RuntimeError. Wrapping consistently across -# tutorials also matches the style of tutorial_01a. +# spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - 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_04a" - BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - PCA_COMPONENTS = 10 - LOG_LEVEL = logging.INFO - - # %% - # Data reading - test_mode = TestTools.running_as_test() + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - log_level = LOG_LEVEL + class_name = "tutorial_04_heart_create_statistical_model" + output_dir = tutorials_dir / "output" / "tutorial_04_heart" + baselines_dir = repo_root / "tests" / "baselines" + + test_mode = TestTools.running_as_test() if test_mode: - pca_components = min(PCA_COMPONENTS, 5) + data_dir = repo_root / "data" / "test" / "KCL-Heart-Model" + pca_components = 5 else: - pca_components = PCA_COMPONENTS + data_dir = repo_root / "data" / "KCL-Heart-Model" + pca_components = 10 + + log_level = logging.INFO + + # Directory setup and data reading output_dir.mkdir(parents=True, exist_ok=True) @@ -86,8 +82,8 @@ reference_mesh = cast(pv.DataSet, pv.read(str(reference_file))) sample_meshes = [cast(pv.DataSet, pv.read(str(path))) for path in sample_files] - # %% # Workflow initialization + workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_meshes, reference_mesh=reference_mesh, @@ -95,11 +91,9 @@ log_level=log_level, ) - # %% # Workflow execution result = workflow.run_workflow() - # %% # Result saving pca_model: dict[str, Any] = result["pca_model"] mean_surface: pv.PolyData = result["pca_mean_surface"] @@ -111,10 +105,11 @@ mean_surface_file = output_dir / "pca_mean_surface.vtp" mean_surface.save(str(mean_surface_file)) + # Testing tt = TestTools( - class_name="tutorial_04a_heart_create_statistical_model", + class_name=class_name, results_dir=output_dir, - baselines_dir=BASELINES_DIR, + baselines_dir=baselines_dir, log_level=log_level, ) @@ -134,35 +129,46 @@ mean_points = np.asarray(mean_surface.points) mode_count = min(2, pca_components, len(components), len(eigenvalues)) + xvfb_started = False try: pv.start_xvfb() + xvfb_started = True except Exception: pass - for mode_idx in range(mode_count): - sigma = float(np.sqrt(eigenvalues[mode_idx])) - mode_offsets = np.asarray(components[mode_idx]).reshape(-1, 3) - - minus_mesh = mean_surface.copy() - minus_mesh.points = mean_points - 2.0 * sigma * mode_offsets - plus_mesh = mean_surface.copy() - plus_mesh.points = mean_points + 2.0 * sigma * mode_offsets - - plotter = pv.Plotter(off_screen=True, window_size=[1200, 500], shape=(1, 3)) - plotter.subplot(0, 0) - plotter.add_mesh(minus_mesh, color="royalblue", opacity=0.9) - plotter.camera_position = "iso" - plotter.subplot(0, 1) - plotter.add_mesh(mean_surface, color="steelblue", opacity=0.9) - plotter.camera_position = "iso" - plotter.subplot(0, 2) - plotter.add_mesh(plus_mesh, color="coral", opacity=0.9) - plotter.camera_position = "iso" - - png_path = output_dir / f"pca_mode_{mode_idx + 1:02d}.png" - plotter.screenshot(str(png_path)) - plotter.close() - screenshots.append(png_path) + try: + for mode_idx in range(mode_count): + sigma = float(np.sqrt(eigenvalues[mode_idx])) + mode_offsets = np.asarray(components[mode_idx]).reshape(-1, 3) + + minus_mesh = mean_surface.copy() + minus_mesh.points = mean_points - 2.0 * sigma * mode_offsets + plus_mesh = mean_surface.copy() + plus_mesh.points = mean_points + 2.0 * sigma * mode_offsets + + plotter = pv.Plotter(off_screen=True, window_size=[1200, 500], shape=(1, 3)) + plotter.subplot(0, 0) + plotter.add_mesh(minus_mesh, color="royalblue", opacity=0.9) + plotter.camera_position = "iso" + plotter.subplot(0, 1) + plotter.add_mesh(mean_surface, color="steelblue", opacity=0.9) + plotter.camera_position = "iso" + plotter.subplot(0, 2) + plotter.add_mesh(plus_mesh, color="coral", opacity=0.9) + plotter.camera_position = "iso" + + png_path = output_dir / f"pca_mode_{mode_idx + 1:02d}.png" + plotter.screenshot(str(png_path)) + plotter.close() + screenshots.append(png_path) + finally: + # Pair start_xvfb with cleanup, guarded like the startup above so + # environments without Xvfb (e.g. Windows, pyvista >= 0.48) are unaffected. + if xvfb_started: + try: + pv.stop_xvfb() + except Exception: + pass tutorial_results = { "pca_model": pca_model, diff --git a/tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py similarity index 69% rename from tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py rename to tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py index 3088a9d..7354759 100644 --- a/tutorials/tutorial_05a_heart_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py @@ -1,19 +1,18 @@ """ -Tutorial 5a: Fit Statistical Shape Model to Patient Data +Tutorial 5: Fit Statistical Shape Model to Patient Data Purpose ------- Fit a generic anatomical template mesh to one or more patient-like surface -meshes. If Tutorial 4a has already written ``pca_model.json``, the workflow uses +meshes. If Tutorial 4 has already written ``pca_model.json``, the workflow uses that model to constrain the fitted shape. Data Required ------------- -Full data: ``data/KCL-Heart-Model`` -Test data: ``data/test/KCL-Heart-Model`` +PCA model: Tutorial 4 output (``output/tutorial_04_heart``) +Patient image: ``data/DirLab-4DCT`` (test: ``data/test/DirLab-4DCT``) """ -# %% # Imports from __future__ import annotations @@ -34,8 +33,7 @@ WorkflowFitStatisticalModelToPatient, ) -# %% -# Data directory specification +# Only run if this script is not imported as a module # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this @@ -43,40 +41,45 @@ # top-level work, that re-import fires the segmenter again and Python's # spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - 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_05a" - BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - 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_academic_license(True) - # SEGMENTATION_METHOD = SegmentHeartSimplewareTrimmedBranches() # Use when available - # and images are contrast-enhanced. - # SEGMENTATION_METHOD = SegmentChestTotalSegmentatorWithContrast() # Use when contrast-enhanced - # images and Simpleware is not available. - LOG_LEVEL = logging.INFO + # Data directory specification + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_05_heart_to_lung_fit_statistical_model_to_patient" + + output_dir = tutorials_dir / "output" / "tutorial_05_heart_to_lung" + baselines_dir = repo_root / "tests" / "baselines" - TEST_MODE = TestTools.running_as_test() + # PCA model + mean surface produced by Tutorial 4. + pca_json = tutorials_dir / "output" / "tutorial_04_heart" / "pca_model.json" + pca_mean_file = ( + tutorials_dir / "output" / "tutorial_04_heart" / "pca_mean_surface.vtp" + ) + + test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + else: + data_dir = repo_root / "data" / "DirLab-4DCT" + patient_image_file = data_dir / "Case1Pack_T70.mha" - data_dir = DATA_DIR - output_dir = OUTPUT_DIR - pca_json = PCA_JSON - pca_mean_file = PCA_MEAN_FILE - patient_image_file = PATIENT_IMAGE_FILE - segmentation_method = SEGMENTATION_METHOD - log_level = LOG_LEVEL + log_level = logging.INFO + + segmentation_method = SegmentChestTotalSegmentator() + 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 images and Simpleware is not available. + + # Directory setup and data reading output_dir.mkdir(parents=True, exist_ok=True) if not pca_mean_file.exists(): raise FileNotFoundError( - f"DirLab-4DCT template not found: {pca_mean_file}\n" - "See data/README.md for download instructions." + f"Tutorial 4 PCA mean surface not found: {pca_mean_file}\n" + "Run Tutorial 4 first (see data/README.md for download instructions)." ) pca_mean = cast(pv.DataSet, pv.read(str(pca_mean_file))) @@ -87,15 +90,13 @@ if not patient_image_file.exists(): raise FileNotFoundError( - f"DirLab-4DCT template not found: {patient_image_file}\n" + f"DirLab-4DCT patient image not found: {patient_image_file}\n" "See data/README.md for download instructions." ) patient_image = itk.imread(str(patient_image_file)) itk.imwrite(patient_image, output_dir / "patient_image.nii.gz") - segmentation_result = segmentation_method.segment( - patient_image, - ) + segmentation_result = segmentation_method.segment(patient_image) patient_labelmap = segmentation_result["labelmap"] itk.imwrite(patient_labelmap, output_dir / "patient_labelmap.nii.gz") @@ -103,13 +104,11 @@ itk.imwrite(heart_labelmap, output_dir / "heart_labelmap.nii.gz") contour_tools = ContourTools() - heart_surface = contour_tools.extract_contours( - labelmap_image=heart_labelmap, - ) + heart_surface = contour_tools.extract_contours(labelmap_image=heart_labelmap) heart_surface.save(output_dir / "heart_surface.vtp") - # %% # Workflow initialization + workflow = WorkflowFitStatisticalModelToPatient( template_model=pca_mean, patient_models=[heart_surface], @@ -126,11 +125,9 @@ use_surface=False, ) - # %% # Workflow execution workflow_results = workflow.process() - # %% # Result saving registered_coefficients = workflow.pca_coefficients if registered_coefficients is not None: @@ -150,7 +147,14 @@ registered_surface = workflow_results["registered_template_model_surface"] registered_surface.save(str(output_dir / "template_surface_registered.vtp")) - # %% + # Testing + TestTools( + class_name=class_name, + results_dir=output_dir, + baselines_dir=baselines_dir, + log_level=log_level, + ) + try: pv.start_xvfb() except Exception: @@ -176,13 +180,6 @@ plotter.close() screenshots.append(after_path) - TestTools( - class_name="tutorial_05a_heart_fit_statistical_model_to_patient", - results_dir=output_dir, - baselines_dir=BASELINES_DIR, - log_level=log_level, - ) - tutorial_results = { "registered_mesh": registered_mesh, "registered_surface": registered_surface, diff --git a/tutorials/tutorial_06_reconstruct_highres_4d_ct.py b/tutorials/tutorial_06_lung_reconstruct_highres_4d_ct.py similarity index 72% rename from tutorials/tutorial_06_reconstruct_highres_4d_ct.py rename to tutorials/tutorial_06_lung_reconstruct_highres_4d_ct.py index ac2c040..f3dfdf5 100644 --- a/tutorials/tutorial_06_reconstruct_highres_4d_ct.py +++ b/tutorials/tutorial_06_lung_reconstruct_highres_4d_ct.py @@ -14,7 +14,6 @@ Test data: ``data/test/DirLab-4DCT/Case1`` """ -# %% # Imports from __future__ import annotations @@ -23,49 +22,52 @@ import itk -from physiotwin4d.register_images_greedy_icon import RegisterImagesGreedyICON -from physiotwin4d.test_tools import TestTools -from physiotwin4d.workflow_reconstruct_highres_4d_ct import ( +from physiotwin4d import ( + RegisterImagesGreedyICON, + TestTools, WorkflowReconstructHighres4DCT, ) +# Only run if this script is not imported as a module + # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this # script in each child; without the __name__ == "__main__" guard around # top-level work, that re-import fires the segmenter again and Python's -# spawn-cascade detector raises RuntimeError. Wrapping consistently across -# tutorials also matches the style of tutorial_01a. +# spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% # Data directory specification - REPO_ROOT = Path(__file__).resolve().parent.parent - TUTORIALS_DIR = Path(__file__).resolve().parent - DATA_DIR = REPO_ROOT / "data" - FULL_DATA_DIR = DATA_DIR / "DirLab-4DCT" - TEST_DATA_DIR = DATA_DIR / "test" / "DirLab-4DCT" + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_06_lung_reconstruct_highres_4d_ct" + + output_dir = tutorials_dir / "output" / "tutorial_06_lung" + baselines_dir = repo_root / "tests" / "baselines" + # .mha files are DirLab-4DCT data already converted to HU by # data/DirLab-4DCT/fix_downloaded_data.py. - CASE_GLOB = "Case1Pack_T??.mha" - OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_06" - BASELINES_DIR = REPO_ROOT / "tests" / "baselines" - LOG_LEVEL = logging.INFO + case_glob = "Case1Pack_T??.mha" - # %% - # Data reading test_mode = TestTools.running_as_test() - - data_dir = TEST_DATA_DIR if test_mode else FULL_DATA_DIR - output_dir = OUTPUT_DIR - log_level = LOG_LEVEL - if test_mode: - number_of_iterations_Greedy = [1, 0] + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + number_of_iterations_greedy = [1, 0] else: - number_of_iterations_Greedy = [30, 15, 7, 3] + data_dir = repo_root / "data" / "DirLab-4DCT" + number_of_iterations_greedy = [30, 15, 7, 3] + + log_level = logging.INFO + + 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 + + # Directory setup and data reading output_dir.mkdir(parents=True, exist_ok=True) - phase_files = sorted(list(data_dir.glob(CASE_GLOB))) + phase_files = sorted(data_dir.glob(case_glob)) if not phase_files: raise FileNotFoundError( f"No DirLab phase images found under {data_dir}.\n" @@ -73,27 +75,22 @@ ) time_series = [itk.imread(str(path)) for path in phase_files] - fixed_image = time_series[0] + reference_image = time_series[0] - # %% # 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, - reference_frame=6, + reference_image=reference_image, + reference_time_frame=6, registration_method=registration_method, log_level=log_level, ) workflow.set_modality("ct") - # %% # Workflow execution result = workflow.process() - # %% # Result saving forward_transform = result["forward_transforms"] inverse_transform = result["inverse_transforms"] @@ -110,17 +107,18 @@ out_path = output_dir / f"reconstructed_frame_{frame_index:03d}_inv.hdf" itk.transformwrite(inverse_transform[frame_index], str(out_path)) + # Testing tt = TestTools( - class_name="tutorial_06_reconstruct_highres_4d_ct", + class_name=class_name, results_dir=output_dir, - baselines_dir=BASELINES_DIR, + baselines_dir=baselines_dir, log_level=log_level, ) screenshots: list[Path] = [] screenshots.append( tt.save_screenshot_image_slice( - fixed_image, + reference_image, "reference_frame.png", axis=0, slice_fraction=0.5, diff --git a/tutorials/tutorial_08cd_byod_fit_model_to_patients.py b/tutorials/tutorial_08_byod_fit_model_to_patients.py similarity index 83% rename from tutorials/tutorial_08cd_byod_fit_model_to_patients.py rename to tutorials/tutorial_08_byod_fit_model_to_patients.py index 6c40f4b..543da63 100644 --- a/tutorials/tutorial_08cd_byod_fit_model_to_patients.py +++ b/tutorials/tutorial_08_byod_fit_model_to_patients.py @@ -1,13 +1,13 @@ """ -Tutorial 8cd: Fit the Cardiac SSM and Propagate It Through Gated Phases +Tutorial 8: Fit the Cardiac SSM and Propagate It Through Gated Phases Purpose ------- -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: +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_09_byod_train_physicsnemo_mgn.py`` / +``tutorial_09_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-7, it does not use the +This is a bring-your-own-data tutorial. Unlike Tutorials 1-6, it does not use the repository ``data/`` directory or a downloadable sample; the path constants below point at a local ``D:/PhysioTwin4D/`` layout. Edit them to match your own data. @@ -36,7 +36,7 @@ * ``D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/`` - PCA model (pca_mean.vtu, pca_model.json) * ``D:/PhysioTwin4D/duke_data/icon_registration/`` - ICON registration weights -Outputs (per patient, under ``OUTPUT_DIR/pm00??/``) +Outputs (per patient, under ``output_dir/pm00??/``) --------------------------------------------------- * ``*_ssm_pca_coefficients.json`` - fitted PCA coefficient vector * ``*_ssm_pca_mesh.vtu`` / ``*_ssm_pca_surface.vtp`` - PCA template before final warp @@ -46,7 +46,6 @@ * ``*_g{TT}_ref_labelmap.nii.gz`` - reference labelmap warped to each phase """ -# %% # Imports from __future__ import annotations @@ -61,11 +60,13 @@ from physiotwin4d import ( ContourTools, RegisterImagesICON, + TestTools, TransformTools, WorkflowFitStatisticalModelToPatient, WorkflowReconstructHighres4DCT, ) -from physiotwin4d.test_tools import TestTools + +# Only run if this script is not imported as a module # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this @@ -73,42 +74,41 @@ # top-level work, that re-import fires the segmenter again and Python's # spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - # %% - # Path configuration (bring-your-own-data: edit for your local layout) - DATA_DIR = Path("D:/PhysioTwin4D/duke_data/gated_nii") - LABELMAP_DIR = Path("D:/PhysioTwin4D/duke_data/simple_ascardio") - SSM_MEAN_MESH_FILE = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") - SSM_MODEL_FILE = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_model.json") - ICON_WEIGHTS_PATH = Path( + # Data directory specification (bring-your-own-data: edit for your local layout) + data_dir = Path("D:/PhysioTwin4D/duke_data/gated_nii") + labelmap_dir = Path("D:/PhysioTwin4D/duke_data/simple_ascardio") + ssm_mean_mesh_file = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_mean.vtu") + ssm_model_file = Path("D:/PhysioTwin4D/kcl-heart-pca/pca-vol-kcl/pca_model.json") + icon_weights_path = Path( "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 9c/9d trainers read from. - OUTPUT_DIR = Path("D:/PhysioTwin4D/duke_data/fitted_kcl_meshes") + # this is also the directory the Tutorial 9 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] + labelmap_interior_object_ids = [1, 2, 3, 4] # Recompute the expensive fit/registration steps (True) or reload cached - # results from OUTPUT_DIR (False). - RECOMPUTE = True - LOG_LEVEL = logging.INFO + # results from output_dir (False). + recompute = True + log_level = logging.INFO - logging.basicConfig(level=LOG_LEVEL) - logger = logging.getLogger("tutorial_08cd_byod_fit_model_to_patients") + class_name = "tutorial_08_byod_fit_model_to_patients" + logging.basicConfig(level=log_level) + logger = logging.getLogger(class_name) # In test mode, limit the run to a single patient to keep it tractable. test_mode = TestTools.running_as_test() + if test_mode: + patient_dirs = sorted(data_dir.glob("pm00??"))[:1] + else: + patient_dirs = sorted(data_dir.glob("pm00??")) - # %% # Load the statistical atlas model - ssm_mean_mesh = pv.read(str(SSM_MEAN_MESH_FILE)) - with SSM_MODEL_FILE.open(encoding="utf-8") as f: + ssm_mean_mesh = pv.read(str(ssm_mean_mesh_file)) + with ssm_model_file.open(encoding="utf-8") as f: ssm_model = json.load(f) - # %% # Discover patients - patient_dirs = sorted(DATA_DIR.glob("pm00??")) - if test_mode: - patient_dirs = patient_dirs[:1] tutorial_results: dict[str, Any] = {"patients": {}} @@ -118,7 +118,7 @@ logger.info("Processing patient %s", patient_id) logger.info("%s", "=" * 48) - patient_output_dir = OUTPUT_DIR / patient_id + patient_output_dir = output_dir / patient_id patient_output_dir.mkdir(parents=True, exist_ok=True) ref_image_files = list(patient_dir.glob("*ref.nii.gz")) @@ -128,9 +128,8 @@ ref_image = itk.imread(str(ref_image_file)) ref_labelmap_file = ref_image_file.name.replace(".nii.gz", "_labelmap.nii.gz") - ref_labelmap = itk.imread(str(LABELMAP_DIR / patient_id / ref_labelmap_file)) + ref_labelmap = itk.imread(str(labelmap_dir / patient_id / ref_labelmap_file)) - # %% # Step 1: fit the statistical model to the reference phase contour_tools = ContourTools() ref_surface = contour_tools.extract_contours(ref_labelmap) @@ -141,14 +140,14 @@ ssm_mesh_path = patient_output_dir / f"{patient_id}_ssm_mesh.vtu" ssm_surface_path = patient_output_dir / f"{patient_id}_ssm_surface.vtp" - if RECOMPUTE: + if recompute: ssm_fit_workflow = WorkflowFitStatisticalModelToPatient( template_model=ssm_mean_mesh, patient_image=ref_image, patient_models=[ref_surface], patient_labelmap=ref_labelmap, - labelmap_interior_object_ids=LABELMAP_INTERIOR_OBJECT_IDS, - log_level=LOG_LEVEL, + labelmap_interior_object_ids=labelmap_interior_object_ids, + log_level=log_level, ) ssm_fit_workflow.set_use_pca_registration( use_pca_registration=True, @@ -189,7 +188,6 @@ ssm_mesh_fitted = pv.read(str(ssm_mesh_path)) ssm_surface_fitted = pv.read(str(ssm_surface_path)) - # %% # Step 2: register every gated phase to the reference gated_files = sorted( file @@ -204,13 +202,13 @@ time_id = gated_file.name.split("_g")[1][:3] time_series_ids.append(time_id) - if RECOMPUTE: + if recompute: icon_registration_method = RegisterImagesICON() - icon_registration_method.set_weights_path(str(ICON_WEIGHTS_PATH)) + icon_registration_method.set_weights_path(str(icon_weights_path)) icon_registration_method.set_number_of_iterations(None) reg_workflow = WorkflowReconstructHighres4DCT( time_series_images=time_series, - fixed_image=ref_image, + reference_image=ref_image, registration_method=icon_registration_method, ) reg_workflow.set_modality("ct") @@ -225,14 +223,13 @@ ) reconstructed_images.append(itk.imread(str(image_path))) - # %% # Step 3: warp the fitted SSM mesh/surface to every gated phase phase_outputs = [] for image_index, image in enumerate(reconstructed_images): time_id = time_series_ids[image_index] logger.info("Patient %s: warping to time point %s", patient_id, time_id) - if RECOMPUTE: + if recompute: image_path = ( patient_output_dir / f"{patient_id}_g{time_id}_warped_ref.mha" ) @@ -255,7 +252,7 @@ ) # Warp the reference labelmap to this phase. Written under - # OUTPUT_DIR (never back into the input labelmap directory). + # output_dir (never back into the input labelmap directory). labelmap = TransformTools().transform_image( ref_labelmap, inv_tfm, image, "nearest" ) diff --git a/tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py similarity index 53% rename from tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py rename to tutorials/tutorial_09_byod_train_physicsnemo_mgn.py index ad24c6f..ec664fe 100644 --- a/tutorials/tutorial_09c_byod_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mgn.py @@ -1,15 +1,15 @@ """ -Tutorial 9c (MGN): Train a PhysicsNeMo MeshGraphNet for cardiac mesh stages. +Tutorial 9 (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 +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +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 +per-time-point SSM surfaces produced by Tutorial 8 +(``tutorial_08_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``. +``tutorial_09_byod_train_physicsnemo_mlp.py``. Why a GNN? ---------- @@ -25,8 +25,8 @@ 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. +Tutorial 8, not at the repository ``data/`` directory. Edit them to match your +own data location. Run Tutorial 8 first. Extra Install Required ---------------------- @@ -35,7 +35,6 @@ pip install "physiotwin4d[physicsnemo]" """ -# %% from __future__ import annotations import json @@ -85,12 +84,13 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ 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" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" + MANIFESTS_DIR = OUTPUT_DIR / "manifests_mgn" + + RESUME_FROM = str(OUTPUT_DIR / "mgn_stage_model_epoch_00100.pt") EPOCHS = 1500 BATCH_SIZE_GRAPHS = 4 # mini-batch measured in (subject, phase) graphs @@ -100,72 +100,65 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ 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"] + TEST_SUBJECTS = ["pm0027"] + VAL_SUBJECTS = [] 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), - ) + """Discover subjects, train a MeshGraphNet, and evaluate the test split.""" + logging.basicConfig(level=LOG_LEVEL) - # 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 - ) + # 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 - return {"training": train_result, "evaluation": eval_outputs} + if len(manifests) < 3: + raise RuntimeError( + f"Found only {len(manifests)} valid subject(s); need at least 3 " + "for a train / val / test split." + ) - # %% - tutorial_results = run_tutorial() + 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, + resume_from=RESUME_FROM, + 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 + ) \ No newline at end of file diff --git a/tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py b/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py similarity index 89% rename from tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py rename to tutorials/tutorial_09_byod_train_physicsnemo_mlp.py index c7dc88c..beaf89e 100644 --- a/tutorials/tutorial_09d_byod_train_physicsnemo_mlp.py +++ b/tutorials/tutorial_09_byod_train_physicsnemo_mlp.py @@ -1,16 +1,16 @@ """ -Tutorial 9d (MLP): Train a PhysicsNeMo fully connected model for cardiac mesh stages. +Tutorial 9 (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 +Second stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +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 +per-time-point SSM surfaces produced by Tutorial 8 +(``tutorial_08_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``) +The companion Tutorial 9 (``tutorial_09_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), @@ -26,8 +26,8 @@ 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. +Tutorial 8, not at the repository ``data/`` directory. Edit them to match your +own data location. Run Tutorial 8 first. Extra Install Required ---------------------- @@ -36,7 +36,6 @@ pip install "physiotwin4d[physicsnemo]" """ -# %% from __future__ import annotations import json @@ -86,12 +85,11 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ 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" + OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mlp" + MANIFESTS_DIR = OUTPUT_DIR / "manifests_mlp" EPOCHS = 10000 BATCH_SIZE_SAMPLES = 32 # (subject, phase) samples per step; points shuffled within @@ -100,8 +98,8 @@ def _write_subject_manifest(subject_dir: Path, manifests_dir: Path) -> Optional[ NUM_LAYERS = 6 # Explicit held-out splits; every other discovered subject is used for training. - TEST_SUBJECTS = ["pm0028"] - VAL_SUBJECTS = ["pm0027"] + TEST_SUBJECTS = ["pm0027"] + VAL_SUBJECTS = [] LOG_LEVEL = logging.INFO def run_tutorial() -> dict[str, Any]: @@ -163,5 +161,4 @@ def run_tutorial() -> dict[str, Any]: return {"training": train_result, "evaluation": eval_outputs} - # %% tutorial_results = run_tutorial() diff --git a/tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py similarity index 84% rename from tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py rename to tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py index 9140ce7..eb65b1a 100644 --- a/tutorials/tutorial_10c_byod_eval_physicsnemo_mgn.py +++ b/tutorials/tutorial_10_byod_eval_physicsnemo_mgn.py @@ -1,22 +1,22 @@ """ -Tutorial 10c (MGN): Predict cardiac stage meshes for a subject with a trained +Tutorial 10 (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 +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +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``). +cardiac surfaces for one subject, loading the model trained by Tutorial 9 +(``tutorial_09_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 +``D:/PhysioTwin4D/`` layout and the Tutorial 9 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 + py tutorial_10_byod_eval_physicsnemo_mgn.py pm0028 --out results/pm0028_mgn + py tutorial_10_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 @@ -40,15 +40,13 @@ 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" +# Tutorial 9 run directory to evaluate (matches that trainer's OUTPUT_DIR). +MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mgn" -DEFAULT_SUBJECT = "pm0028" -DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10c" / DEFAULT_SUBJECT +DEFAULT_SUBJECT = "pm0027" +DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10_byod_mgn" / DEFAULT_SUBJECT def _gating_stage_from_filename(mesh_file: Path) -> float: @@ -130,6 +128,5 @@ def main() -> None: 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_10_byod_eval_physicsnemo_mlp.py similarity index 84% rename from tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py rename to tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py index 5e644dc..f7d1fb0 100644 --- a/tutorials/tutorial_10d_byod_eval_physicsnemo_mlp.py +++ b/tutorials/tutorial_10_byod_eval_physicsnemo_mlp.py @@ -1,22 +1,22 @@ """ -Tutorial 10d (MLP): Predict cardiac stage meshes for a subject with a trained +Tutorial 10 (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 +Final stage of the cardiac 4D deep-learning pipeline (Tutorials 8 -> 9 -> 10). +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``). +cardiac surfaces for one subject, loading the model trained by Tutorial 9 +(``tutorial_09_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 +``D:/PhysioTwin4D/`` layout and the Tutorial 9 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 + py tutorial_10_byod_eval_physicsnemo_mlp.py pm0028 --out results/pm0028_mlp + py tutorial_10_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 @@ -40,15 +40,13 @@ 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" +# Tutorial 9 run directory to evaluate (matches that trainer's OUTPUT_DIR). +MODEL_DIR = TUTORIALS_DIR / "output" / "tutorial_09_byod_mlp" -DEFAULT_SUBJECT = "pm0028" -DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10d" / DEFAULT_SUBJECT +DEFAULT_SUBJECT = "pm0027" +DEFAULT_OUT_DIR = MODEL_DIR / "tutorial_10_byod_mlp" / DEFAULT_SUBJECT def _gating_stage_from_filename(mesh_file: Path) -> float: @@ -130,6 +128,5 @@ def main() -> None: if len(sys.argv) > 1: main() else: - # %% # Tutorial / test entry point (no CLI arguments) tutorial_results = run_tutorial()