From 92dff6aa1273eb77d9e77eebb0d027ce3b035223 Mon Sep 17 00:00:00 2001 From: "Stephen R. Aylward" Date: Thu, 14 May 2026 10:28:27 -0400 Subject: [PATCH 1/4] docs: add i4h BYOD DICOM & VTK to USD tutorial page (3D + 4D) --- docs/cli_scripts/byod_tutorials.rst | 240 ++++++++++++++++++++++++++++ 1 file changed, 240 insertions(+) create mode 100644 docs/cli_scripts/byod_tutorials.rst diff --git a/docs/cli_scripts/byod_tutorials.rst b/docs/cli_scripts/byod_tutorials.rst new file mode 100644 index 0000000..eb67122 --- /dev/null +++ b/docs/cli_scripts/byod_tutorials.rst @@ -0,0 +1,240 @@ +.. _byod_tutorials: + +Bring Your Own Data — DICOM & VTK to USD +========================================= + +PhysioMotion4D lets you convert your own medical imaging data — whether +DICOM-derived NIfTI volumes or VTK surface meshes — into OpenUSD for +interactive visualization in NVIDIA Omniverse. Both 3D (single +volume/mesh) and 4D (time-series) inputs are supported. The CLI and +Python API are **identical** for 3D and 4D inputs; the only difference +is how many files you pass in. + +.. note:: + + PhysioMotion4D is a research tool and has **not** been validated for + clinical use. Outputs must not be used for diagnostic or therapeutic + decisions without independent validation. + +Installation +------------ + +Install the package with CUDA support (recommended for GPU acceleration) +or the CPU-only variant: + +.. code-block:: bash + + # Recommended — CUDA-enabled + pip install physiomotion4d[cuda] + + # CPU-only + pip install physiomotion4d + +Verify that both relevant CLI entry-points are available after installation: + +.. code-block:: bash + + physiomotion4d-convert-heart-gated-ct-to-usd --help + physiomotion4d-convert-vtk-to-usd --help + +See :doc:`/installation` for prerequisites, CUDA version requirements, and +source-based installation. + +DICOM to USD +------------ + +Raw DICOM images must first be converted to NIfTI with a tool such as +`dcm2niix `_ before being passed +to PhysioMotion4D. + +3D — Single Volume +~~~~~~~~~~~~~~~~~~ + +Pass a single ``.nii.gz`` file to produce a static USD scene. + +**CLI:** + +.. code-block:: bash + + physiomotion4d-convert-heart-gated-ct-to-usd \ + patient_ct.nii.gz \ + --output patient_heart.usd + +**Python API:** + +.. code-block:: python + + import physiomotion4d as pm4d + + wf = pm4d.WorkflowConvertHeartGatedCTToUSD() + wf.input_image = "patient_ct.nii.gz" + wf.output_file = "patient_heart.usd" + wf.run_workflow() + +.. note:: + + If your source data is raw DICOM, run ``dcm2niix -z y -o output_dir + dicom_dir/`` first to produce the ``.nii.gz`` input file expected by + this command. + +4D — Gated CT Time Series +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pass multiple per-phase volumes (glob or explicit list) to produce an +animated USD scene. Use ``--fps`` to control playback rate and +``--reference-frame`` to choose the registration anchor phase (0-indexed). + +**CLI:** + +.. code-block:: bash + + physiomotion4d-convert-heart-gated-ct-to-usd \ + phase_*.nii.gz \ + --output heart_animated.usd \ + --fps 25 \ + --reference-frame 0 + +**Python API:** + +.. code-block:: python + + import glob + import physiomotion4d as pm4d + + wf = pm4d.WorkflowConvertHeartGatedCTToUSD() + wf.input_images = sorted(glob.glob("phase_*.nii.gz")) + wf.output_file = "heart_animated.usd" + wf.fps = 25 + wf.reference_frame = 0 + wf.run_workflow() + +The resulting USD file contains a time-sampled mesh sequence that plays +back when you press **Play** in Omniverse USD Composer. + +VTK to USD +---------- + +3D — Single Mesh +~~~~~~~~~~~~~~~~ + +Pass a single ``.vtp`` file. Use ``--appearance`` to control material +style and ``--no-split`` to skip the default connected-component split. + +**CLI:** + +.. code-block:: bash + + # Default — split by connected component, anatomy material + physiomotion4d-convert-vtk-to-usd heart.vtp \ + --output heart.usd \ + --appearance anatomy \ + --anatomy-type heart + + # Solid colour, no splitting + physiomotion4d-convert-vtk-to-usd mesh.vtp \ + --output mesh_red.usd \ + --appearance solid \ + --color 0.8 0.1 0.1 \ + --no-split + +**Python API:** + +.. code-block:: python + + import physiomotion4d as pm4d + + wf = pm4d.WorkflowConvertVTKToUSD() + wf.input_files = ["heart.vtp"] + wf.output_file = "heart.usd" + wf.appearance = "anatomy" + wf.anatomy_type = "heart" + wf.run() + +4D — Mesh Time Series +~~~~~~~~~~~~~~~~~~~~~ + +Pass a glob of per-frame files. The ``--fps`` flag controls playback +rate. For scalar colormaps, combine ``--primvar``, ``--cmap``, and +``--intensity-range``. + +**CLI:** + +.. code-block:: bash + + # Animated mesh sequence + physiomotion4d-convert-vtk-to-usd frame_*.vtp \ + --output heart_animation.usd \ + --fps 30 + + # Animated with scalar colormap (e.g. wall stress) + physiomotion4d-convert-vtk-to-usd frame_*.vtk \ + --output stress_animation.usd \ + --fps 30 \ + --appearance colormap \ + --primvar vtk_point_stress_c0 \ + --cmap viridis \ + --intensity-range 0 500 + +**Python API:** + +.. code-block:: python + + import glob + import physiomotion4d as pm4d + + wf = pm4d.WorkflowConvertVTKToUSD() + wf.input_files = sorted(glob.glob("frame_*.vtk")) + wf.output_file = "stress_animation.usd" + wf.fps = 30 + wf.appearance = "colormap" + wf.primvar = "vtk_point_stress_c0" + wf.cmap = "viridis" + wf.intensity_range = (0, 500) + wf.run() + +**Lower-level in-memory conversion with** ``ConvertVTKToUSD``**:** + +For programmatic pipelines where meshes are already in memory, use the +lower-level :class:`physiomotion4d.ConvertVTKToUSD` class directly: + +.. code-block:: python + + import pyvista as pv + import physiomotion4d as pm4d + + # Load or construct meshes in memory + meshes = [pv.read(f"frame_{i:04d}.vtp") for i in range(10)] + + converter = pm4d.ConvertVTKToUSD(output_file="output.usd", fps=30) + for i, mesh in enumerate(meshes): + converter.add_frame(mesh, frame_index=i) + converter.write() + +Viewing Results +--------------- + +**Quick preview with PyVista (no Omniverse required):** + +.. code-block:: python + + import pyvista as pv + + mesh = pv.read("output.usd") + mesh.plot() + +**In NVIDIA Omniverse:** + +Open **Omniverse USD Composer**, drag your ``.usd`` file onto the +viewport, then press **Play** (spacebar) to watch the animation. For +4D cardiac data, use the **Timeline** panel to scrub through phases. + +See Also +-------- + +- :doc:`/installation` +- :doc:`/quickstart` +- :doc:`/cli_scripts/heart_gated_ct` +- :doc:`/cli_scripts/vtk_to_usd` +- :doc:`/api/workflows` +- :doc:`/examples` +- :doc:`/troubleshooting` From 1df925cfe42614fdb092f1e036915ae02208a406 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sat, 16 May 2026 10:33:16 -0400 Subject: [PATCH 2/4] BUG: Fix cuda installation - [cuda13] --- docs/cli_scripts/byod_tutorials.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cli_scripts/byod_tutorials.rst b/docs/cli_scripts/byod_tutorials.rst index eb67122..fa6a122 100644 --- a/docs/cli_scripts/byod_tutorials.rst +++ b/docs/cli_scripts/byod_tutorials.rst @@ -25,7 +25,7 @@ or the CPU-only variant: .. code-block:: bash # Recommended — CUDA-enabled - pip install physiomotion4d[cuda] + pip install physiomotion4d[cuda13] # CPU-only pip install physiomotion4d From 74f527d9265b2394c17b62ca9939c1aa3882613a Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sun, 17 May 2026 14:36:09 -0400 Subject: [PATCH 3/4] =?UTF-8?q?ENH:=20Generalize=20CT=E2=86=92Image=20nami?= =?UTF-8?q?ng,=20expose=20backend=20selectors,=20add=20ICON=20finetune?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename WorkflowConvertCTToVTK → WorkflowConvertImageToVTK and its CLI (physiomotion4d-convert-ct-to-vtk → -convert-image-to-vtk); the workflow is image-modality-agnostic, not CT-specific. - Expose --segmentation-method and --registration-method on the ImageToVTK and ImageToUSD CLIs. Backend identifiers: 'ChestTotalSegmentator' / 'HeartSimpleware' and 'ANTS' / 'ICON' (breaking: replaces lowercase total_segmentator / simpleware_heart / ants / icon). - Add segmentation_method to WorkflowConvertImageToUSD (previously hardcoded to TotalSegmentator). - Rename SegmentHeartSimpleware.set_trim_mask_to_essentials → set_trim_branches (and trim_mask_to_essentials → trim_branches), with docs noting that trimming reduces inter-subject vessel-extent variability for AI-Ready / Sim-Ready model fitting, consistent with the KCL Heart example data. - Add RegisterImagesICON.finetune() for cohort fine-tuning of UniGradICON weights with optional mask pairs, persistent weight updates, and checkpoint export. - Propagate new identifiers through tutorials, experiments, docs, and regenerate API_MAP.md. - Add torch to mypy ignore_missing_imports override list. - Refresh CLAUDE.md with role, priorities, and behavior guidelines. --- CLAUDE.md | 39 +++- README.md | 4 +- docs/API_MAP.md | 55 +++--- docs/api/cli/convert_ct_to_vtk.rst | 10 - docs/api/cli/convert_image_to_vtk.rst | 10 + docs/api/cli/index.rst | 2 +- docs/api/segmentation/totalsegmentator.rst | 2 +- docs/api/workflows.rst | 19 +- docs/architecture.rst | 6 +- docs/cli_scripts/best_practices.rst | 2 +- docs/cli_scripts/heart_gated_ct.rst | 8 +- docs/cli_scripts/overview.rst | 4 +- docs/developer/architecture.rst | 4 +- docs/developer/workflows.rst | 6 +- docs/examples.rst | 6 +- docs/quickstart.rst | 8 +- docs/troubleshooting.rst | 4 +- docs/tutorials.rst | 4 +- .../simpleware_heart_segmentation.py | 2 +- .../heart_model_to_patient-CHOPValve.py | 2 +- pyproject.toml | 4 +- src/physiomotion4d/__init__.py | 4 +- .../cli/convert_image_to_usd.py | 21 +- ...t_ct_to_vtk.py => convert_image_to_vtk.py} | 32 ++-- src/physiomotion4d/register_images_icon.py | 180 ++++++++++++++++-- .../segment_heart_simpleware.py | 29 ++- .../workflow_convert_image_to_usd.py | 48 ++++- ...tk.py => workflow_convert_image_to_vtk.py} | 44 ++--- ...rkflow_fit_statistical_model_to_patient.py | 8 +- tests/test_cli_smoke.py | 2 +- tutorials/README.md | 2 +- .../tutorial_01_heart_gated_ct_to_usd.py | 6 +- tutorials/tutorial_02_ct_to_vtk.py | 10 +- tutorials/tutorial_07_dirlab_pca_model.py | 6 +- 34 files changed, 414 insertions(+), 179 deletions(-) delete mode 100644 docs/api/cli/convert_ct_to_vtk.rst create mode 100644 docs/api/cli/convert_image_to_vtk.rst rename src/physiomotion4d/cli/{convert_ct_to_vtk.py => convert_image_to_vtk.py} (88%) rename src/physiomotion4d/{workflow_convert_ct_to_vtk.py => workflow_convert_image_to_vtk.py} (92%) diff --git a/CLAUDE.md b/CLAUDE.md index 7119223..aca3523 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,6 +5,22 @@ Project guidance for Claude Code in this repository. Codex and other AI agents should use `AGENTS.md` as the primary shared instructions file. Claude-specific behavior and slash-command usage remain here. +## Role: +We are developing open-source code for scientific AI libraries. Leverage GPU-accelerated methods when appropriate. + +## Priorities (Ordered) +1) accuracy +2) clarity/maintainability/simplicity +3) consistency with the rest of the platform and open source standards +4) documentation +5) testing + +## Behavior Guidelines +1) Don't assume. Don't hide confuction. Surface tradeoffs. +2) Minimum code that solves the problem. Nothing speculative. +3) Touch only what you must. Clean up only your own mess. +4) Define success criteria. Loop until verified. + ## Commands **Python launcher:** Use `py` on this Windows system (not `python`). @@ -47,8 +63,6 @@ py -m pytest tests/ --create-baselines ## Architecture -Pipeline: `4D CT → Segmentation → Registration → Contour Extraction → USD Export` - All classes inherit from `PhysioMotion4DBase` (`physiomotion4d_base.py`), which provides a shared logger. Use `self.log_info()`, `self.log_debug()` — never `print()`. @@ -58,12 +72,13 @@ Regenerate it after any public API change: `py utils/generate_api_map.py` **Key data conventions:** - Images: `itk.Image`, axes X, Y, Z [, T] in LPS world space (ITK's native frame; `itk.imread` normalizes DICOM, NIfTI, MHA, and NRRD inputs to LPS) + stored using itk.imwrite with compression=True - 4D time series: shape `(X, Y, Z, T)` — never silently squeeze or permute axes - Surfaces: `pv.PolyData` in LPS (inherited from the source `itk.Image` via `itk.vtk_image_from_image`); converted to USD right-handed Y-up only at USD export by `vtk_to_usd.lps_points_to_usd` (USD +X=Left, +Y=Superior, +Z=Anterior) - Masks: ITK images with integer labels; consistent anatomy group IDs across all segmenters -- Transforms: ITK composite transforms stored in `.hdf` files +- Transforms: ITK composite transforms stored in `.hdf` files with compression - State axis order and shape explicitly in every docstring and comment that touches arrays ## Testing @@ -71,18 +86,23 @@ Regenerate it after any public API change: `py utils/generate_api_map.py` - Baselines in `tests/baselines/` via Git LFS — run `git lfs pull` after cloning - `tests/conftest.py`: session-scoped fixtures chaining download → convert → segment → register - `src/physiomotion4d/test_tools.py`: baseline comparison utilities (`TestTools`, etc.) -- Markers: `slow`, `requires_gpu`, `requires_data`, `experiment` (skipped by default) -- Prefer synthetic `itk.Image` / `pv.PolyData` over real data; keep volumes ≤64 voxels/side +- Markers: `slow`, `requires_gpu`, `requires_data`, `experiment`, `tutorial` +- Prefer images from `ROOT/data/test/slicer_heart_small` for tests +- Prefer storing results in subdirs `./results/` ## Working Process Before editing any code: 1. Read the relevant source file(s) in full. 2. Summarize current behavior in 2–4 sentences. -3. Propose a numbered plan; confirm before implementing non-trivial changes. -4. Implement in small, reviewable diffs. -5. Update docstrings and tests for every changed public method. -6. Call out breaking changes explicitly. +3. Identify success criteria / metrics +4. Refer to *_tools.py files for commonly used routines +5. Refer to workspace/reference_code (when available) for third-party libraries +6. Propose a numbered plan; confirm before implementing non-trivial changes. +7. Follow the behavior guidelines given above. +8. Implement in small, reviewable diffs. +9. Update docstrings and tests for every changed public method. +10. Call out breaking changes explicitly. Breaking changes are acceptable. Backward-compatibility shims are not. @@ -113,3 +133,4 @@ Document via docstrings and inline comments. - `Optional[X]` not `X | None` (ruff `UP007` suppressed) - Breaking changes are acceptable — backward compatibility is not a priority - Max line length: 88 characters +- Follow behavior guidelines. diff --git a/README.md b/README.md index 8a165a3..7a87fff 100644 --- a/README.md +++ b/README.md @@ -196,7 +196,7 @@ a CUDA-capable GPU are required for practical runtime. python -c "from physiomotion4d import DataDownloadTools; DataDownloadTools.DownloadSlicerHeartCTData('data/test')" physiomotion4d-convert-image-to-usd data/test/TruncalValve_4DCT.seq.nrrd \ - --registration-method ants \ + --registration-method ANTS \ --output-dir output/quickstart \ --project-name slicer_heart_quickstart ``` @@ -286,7 +286,7 @@ processor = WorkflowConvertImageToUSD( contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", - registration_method='icon' # or 'ants' + registration_method='ICON' # or 'ANTS' ) # Run complete workflow diff --git a/docs/API_MAP.md b/docs/API_MAP.md index 411e2cf..b3493a4 100644 --- a/docs/API_MAP.md +++ b/docs/API_MAP.md @@ -62,10 +62,6 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def group_for_id(self, label_id)` (line 132): Return the group containing *label_id*; :data:`OTHER_GROUP` if absent. - `def fill_other_group(self, id_range=range(1, 256), name_template='other_{id}')` (line 139): Populate the ``other`` group with any ids not already claimed. -## src/physiomotion4d/cli/convert_ct_to_vtk.py - -- `def main()` (line 26): CLI entry point for CT to VTK conversion. - ## src/physiomotion4d/cli/convert_image_4d_to_3d.py - `def main()` (line 20): CLI entry point for 4D-to-3D image conversion. @@ -74,6 +70,10 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def main()` (line 14): Command-line interface for the Image-to-USD workflow. +## src/physiomotion4d/cli/convert_image_to_vtk.py + +- `def main()` (line 26): CLI entry point for image to VTK conversion. + ## src/physiomotion4d/cli/convert_vtk_to_usd.py - `def main()` (line 22): Command-line interface for VTK to USD conversion. @@ -202,14 +202,15 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ ## src/physiomotion4d/register_images_icon.py -- **class RegisterImagesICON** (line 25): ICON-based deformable image registration implementation. - - `def __init__(self, log_level=logging.INFO)` (line 61): Initialize the ICON image registration class. - - `def set_weights_path(self, weights_path)` (line 79): Set a custom weights file for the uniGradICON network. - - `def set_number_of_iterations(self, number_of_iterations)` (line 93): Set the number of iterations for ICON registration. - - `def set_multi_modality(self, enable)` (line 101): Enable or disable multi-modality registration. - - `def set_mass_preservation(self, enable)` (line 118): Enable or disable mass preservation constraint. - - `def preprocess(self, image, modality='ct')` (line 135): Preprocess the image for ICON registration. - - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 155): Register moving image to fixed image using ICON registration algorithm. +- **class RegisterImagesICON** (line 32): ICON-based deformable image registration implementation. + - `def __init__(self, log_level=logging.INFO)` (line 68): Initialize the ICON image registration class. + - `def set_weights_path(self, weights_path)` (line 86): Set a custom weights file for the uniGradICON network. + - `def set_number_of_iterations(self, number_of_iterations)` (line 100): Set the number of iterations for ICON registration. + - `def set_multi_modality(self, enable)` (line 108): Enable or disable multi-modality registration. + - `def set_mass_preservation(self, enable)` (line 125): Enable or disable mass preservation constraint. + - `def preprocess(self, image, modality='ct')` (line 142): Preprocess the image for ICON registration. + - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 162): Register moving image to fixed image using ICON registration algorithm. + - `def finetune(self, image_pairs, output_model_filename, mask_pairs=None, epochs=1, learning_rate=DEFAULT_FINETUNE_LEARNING_RATE)` (line 336): Fine-tune the ICON network on a cohort of image pairs. ## src/physiomotion4d/register_models_distance_maps.py @@ -285,11 +286,11 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - **class SegmentHeartSimpleware** (line 23): Heart CT segmentation using Simpleware Medical's ASCardio module. - `def __init__(self, log_level=logging.INFO)` (line 58): Initialize the Simpleware Medical based heart segmentation. - - `def set_trim_mask_to_essentials(self, trim_mask)` (line 118): Set whether to trim mask to common and critical structures. - - `def set_simpleware_executable_path(self, path)` (line 126): Set the path to the Simpleware Medical console executable. - - `def segmentation_method(self, preprocessed_image)` (line 139): Run Simpleware Medical ASCardio segmentation on the preprocessed image. - - `def get_landmarks(self)` (line 339): Get the landmarks. - - `def trim_mask_to_essentials(self, labelmap_image)` (line 343): Trim mask to essentials. + - `def set_trim_branches(self, trim_branches)` (line 118): Enable trimming of pulmonary and great-vessel branches. + - `def set_simpleware_executable_path(self, path)` (line 133): Set the path to the Simpleware Medical console executable. + - `def segmentation_method(self, preprocessed_image)` (line 146): Run Simpleware Medical ASCardio segmentation on the preprocessed image. + - `def get_landmarks(self)` (line 346): Get the landmarks. + - `def trim_branches(self, labelmap_image)` (line 350): Trim pulmonary and great-vessel branches back to the cardiac region. ## src/physiomotion4d/test_tools.py @@ -425,22 +426,22 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def read_vtk_file(filename, extract_surface=True)` (line 581): Auto-detect VTK file format and read appropriately. - `def validate_time_series_topology(mesh_data_sequence, filenames=None)` (line 609): Validate topology consistency across a time series of meshes. -## src/physiomotion4d/workflow_convert_ct_to_vtk.py +## src/physiomotion4d/workflow_convert_image_to_usd.py + +- **class WorkflowConvertImageToUSD** (line 41): Complete workflow for converting 4D CT images to dynamic USD models. + - `def __init__(self, input_filenames, contrast_enhanced, output_directory, project_name, reference_image_filename=None, number_of_registration_iterations=1, segmentation_method='ChestTotalSegmentator', registration_method='ICON', log_level=logging.INFO, save_registered_images=True, save_registration_transforms=True, save_labelmaps=True)` (line 49): Initialize the image-to-USD workflow. + - `def process(self)` (line 214): Execute the complete workflow from 4D CT to dynamic USD models. + +## src/physiomotion4d/workflow_convert_image_to_vtk.py -- **class WorkflowConvertCTToVTK** (line 58): Segment a CT image and produce per-anatomy-group VTK surfaces and meshes. - - `def __init__(self, segmentation_method='total_segmentator', log_level=logging.INFO)` (line 98): Initialize the workflow. +- **class WorkflowConvertImageToVTK** (line 58): Segment a CT image and produce per-anatomy-group VTK surfaces and meshes. + - `def __init__(self, segmentation_method='ChestTotalSegmentator', log_level=logging.INFO)` (line 98): Initialize the workflow. - `def run_workflow(self, input_image, contrast_enhanced_study=False, anatomy_groups=None)` (line 241): Segment the CT image and extract per-anatomy-group VTK objects. - `def save_surfaces(surfaces, output_dir, prefix='')` (line 344): Save each group surface to its own VTP file. - `def save_meshes(meshes, output_dir, prefix='')` (line 371): Save each group voxel mesh to its own VTU file. - `def save_combined_surface(surfaces, output_dir, prefix='')` (line 397): Merge all group surfaces into a single VTP file. - `def save_combined_mesh(meshes, output_dir, prefix='')` (line 432): Merge all group meshes into a single VTU file. -## src/physiomotion4d/workflow_convert_image_to_usd.py - -- **class WorkflowConvertImageToUSD** (line 30): Complete workflow for converting 4D CT images to dynamic USD models. - - `def __init__(self, input_filenames, contrast_enhanced, output_directory, project_name, reference_image_filename=None, number_of_registration_iterations=1, registration_method='icon', log_level=logging.INFO, save_registered_images=True, save_registration_transforms=True, save_labelmaps=True)` (line 38): Initialize the image-to-USD workflow. - - `def process(self)` (line 184): Execute the complete workflow from 4D CT to dynamic USD models. - ## src/physiomotion4d/workflow_convert_vtk_to_usd.py - **class WorkflowConvertVTKToUSD** (line 23): Workflow to convert one or more VTK files to USD with configurable @@ -458,7 +459,7 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ ## src/physiomotion4d/workflow_fit_statistical_model_to_patient.py - **class WorkflowFitStatisticalModelToPatient** (line 56): Register anatomical models using multi-stage ICP, mask-based, and image-based - - `def __init__(self, template_model, patient_models=None, patient_image=None, segmentation_method='simpleware_heart', log_level=logging.INFO)` (line 135): Initialize the model-to-image-and-model registration pipeline. + - `def __init__(self, template_model, patient_models=None, patient_image=None, segmentation_method='HeartSimpleware', log_level=logging.INFO)` (line 135): Initialize the model-to-image-and-model registration pipeline. - `def set_mask_dilation_mm(self, mask_dilation_mm)` (line 356): Set mask dilation amount for auto-generated masks. - `def set_roi_dilation_mm(self, roi_dilation_mm)` (line 365): Set ROI mask dilation amount. - `def set_use_pca_registration(self, use_pca_registration, pca_model=None, pca_number_of_modes=0, pca_uses_surface=True)` (line 374): Set whether to use PCA-based registration and provide the PCA model. diff --git a/docs/api/cli/convert_ct_to_vtk.rst b/docs/api/cli/convert_ct_to_vtk.rst deleted file mode 100644 index e910d8d..0000000 --- a/docs/api/cli/convert_ct_to_vtk.rst +++ /dev/null @@ -1,10 +0,0 @@ -======================= -convert_ct_to_vtk (CLI) -======================= - -.. automodule:: physiomotion4d.cli.convert_ct_to_vtk - :members: - :undoc-members: - :show-inheritance: - -See :doc:`../../cli_scripts/overview` for usage. diff --git a/docs/api/cli/convert_image_to_vtk.rst b/docs/api/cli/convert_image_to_vtk.rst new file mode 100644 index 0000000..c132590 --- /dev/null +++ b/docs/api/cli/convert_image_to_vtk.rst @@ -0,0 +1,10 @@ +========================== +convert_image_to_vtk (CLI) +========================== + +.. automodule:: physiomotion4d.cli.convert_image_to_vtk + :members: + :undoc-members: + :show-inheritance: + +See :doc:`../../cli_scripts/overview` for usage. diff --git a/docs/api/cli/index.rst b/docs/api/cli/index.rst index aa885c3..1e57083 100644 --- a/docs/api/cli/index.rst +++ b/docs/api/cli/index.rst @@ -19,7 +19,7 @@ Module Index. :maxdepth: 1 convert_image_to_usd - convert_ct_to_vtk + convert_image_to_vtk convert_vtk_to_usd create_statistical_model fit_statistical_model_to_patient diff --git a/docs/api/segmentation/totalsegmentator.rst b/docs/api/segmentation/totalsegmentator.rst index 2b11776..4f38e21 100644 --- a/docs/api/segmentation/totalsegmentator.rst +++ b/docs/api/segmentation/totalsegmentator.rst @@ -67,7 +67,7 @@ Operational Notes TotalSegmentator model inference may download model assets and can be slow on a CPU-only environment. For repeatable workflows, prefer the tutorial scripts or -the ``physiomotion4d-convert-ct-to-vtk`` CLI. +the ``physiomotion4d-convert-image-to-vtk`` CLI. See Also ======== diff --git a/docs/api/workflows.rst b/docs/api/workflows.rst index 0fe4a36..6b20255 100644 --- a/docs/api/workflows.rst +++ b/docs/api/workflows.rst @@ -3,7 +3,7 @@ Workflow Classes ================ .. module:: physiomotion4d.workflow_convert_image_to_usd -.. module:: physiomotion4d.workflow_convert_ct_to_vtk +.. module:: physiomotion4d.workflow_convert_image_to_vtk .. module:: physiomotion4d.workflow_convert_vtk_to_usd .. module:: physiomotion4d.workflow_create_statistical_model .. module:: physiomotion4d.workflow_fit_statistical_model_to_patient @@ -26,7 +26,7 @@ Available Workflows - Typical use * - :class:`WorkflowConvertImageToUSD` - Convert a 4D cardiac CT sequence into animated USD anatomy. - * - :class:`WorkflowConvertCTToVTK` + * - :class:`WorkflowConvertImageToVTK` - Segment one CT image and export anatomy-group VTK surfaces and meshes. * - :class:`WorkflowConvertVTKToUSD` - Convert VTK/VTP/VTU meshes or time series into USD. @@ -55,15 +55,16 @@ Convert Image to USD contrast_enhanced=True, output_directory="./results", project_name="patient_001", - registration_method="ants", + segmentation_method="ChestTotalSegmentator", + registration_method="ANTS", ) final_usd = workflow.process() -CT to VTK -========= +Image to VTK +============ -.. autoclass:: WorkflowConvertCTToVTK +.. autoclass:: WorkflowConvertImageToVTK :members: :undoc-members: :show-inheritance: @@ -72,17 +73,17 @@ CT to VTK import itk - from physiomotion4d import WorkflowConvertCTToVTK + from physiomotion4d import WorkflowConvertImageToVTK image = itk.imread("chest_ct.nii.gz") - workflow = WorkflowConvertCTToVTK(segmentation_method="total_segmentator") + workflow = WorkflowConvertImageToVTK(segmentation_method="ChestTotalSegmentator") result = workflow.run_workflow( input_image=image, contrast_enhanced_study=True, anatomy_groups=["heart", "major_vessels"], ) - WorkflowConvertCTToVTK.save_combined_surface( + WorkflowConvertImageToVTK.save_combined_surface( result["surfaces"], "./output", prefix="patient01", diff --git a/docs/architecture.rst b/docs/architecture.rst index 93f4a40..7524cd8 100644 --- a/docs/architecture.rst +++ b/docs/architecture.rst @@ -33,7 +33,7 @@ Data Flow ContourTools + TransformTools | v - WorkflowConvertCTToVTK / ConvertVTKToUSD / WorkflowConvertVTKToUSD + WorkflowConvertImageToVTK / ConvertVTKToUSD / WorkflowConvertVTKToUSD | v OpenUSD assets for NVIDIA Omniverse @@ -45,7 +45,7 @@ Primary Workflows Converts a 4D cardiac CT file or 3D CT time series into registered anatomy contours and painted animated USD files. -``WorkflowConvertCTToVTK`` +``WorkflowConvertImageToVTK`` Segments a 3D CT image and exports anatomy groups as VTK surfaces and voxel meshes. @@ -89,7 +89,7 @@ The installed CLI commands in ``pyproject.toml`` are thin wrappers around the workflow classes. They are the preferred examples for executable API usage: * ``physiomotion4d-convert-image-to-usd`` -* ``physiomotion4d-convert-ct-to-vtk`` +* ``physiomotion4d-convert-image-to-vtk`` * ``physiomotion4d-create-statistical-model`` * ``physiomotion4d-fit-statistical-model-to-patient`` * ``physiomotion4d-convert-vtk-to-usd`` diff --git a/docs/cli_scripts/best_practices.rst b/docs/cli_scripts/best_practices.rst index fd01ed6..532839e 100644 --- a/docs/cli_scripts/best_practices.rst +++ b/docs/cli_scripts/best_practices.rst @@ -171,7 +171,7 @@ Registration Failure **Solutions**: * Select better reference image (less motion, better quality) * Increase ``--registration-iterations`` - * Try alternative registration method (``--registration-method ants``) + * Try alternative registration method (``--registration-method ANTS``) * Verify phases are temporally ordered correctly Memory Errors diff --git a/docs/cli_scripts/heart_gated_ct.rst b/docs/cli_scripts/heart_gated_ct.rst index 059c2d8..d42a975 100644 --- a/docs/cli_scripts/heart_gated_ct.rst +++ b/docs/cli_scripts/heart_gated_ct.rst @@ -111,8 +111,8 @@ Optional Arguments - 1 - Number of registration refinement iterations * - ``--registration-method`` - - ``icon`` - - Registration method: ``icon`` or ``ants`` + - ``ICON`` + - Registration method: ``ICON`` or ``ANTS`` Processing Pipeline =================== @@ -222,7 +222,7 @@ With ANTs Registration physiomotion4d-convert-image-to-usd cardiac.nrrd \ --contrast \ - --registration-method ants \ + --registration-method ANTS \ --registration-iterations 50 Best Practices @@ -289,7 +289,7 @@ Registration Failures **Registration not converging** * Try different reference image (better quality phase) * Increase ``--registration-iterations`` - * Switch registration method (``--registration-method ants``) + * Switch registration method (``--registration-method ANTS``) **Excessive deformation** * Verify sufficient temporal overlap between phases diff --git a/docs/cli_scripts/overview.rst b/docs/cli_scripts/overview.rst index 60058b4..2ed59a4 100644 --- a/docs/cli_scripts/overview.rst +++ b/docs/cli_scripts/overview.rst @@ -51,8 +51,8 @@ Current Scripts - Description * - :doc:`heart_gated_ct` - Process cardiac gated CT to animated heart models with physiological motion - * - ``physiomotion4d-convert-ct-to-vtk`` - - Segment one CT image and export anatomy-group VTK surfaces and meshes + * - ``physiomotion4d-convert-image-to-vtk`` + - Segment one 3D image and export anatomy-group VTK surfaces and meshes * - ``physiomotion4d-convert-image-4d-to-3d`` - Split a 4D medical image into a 3D time series using ITK readers * - :doc:`create_statistical_model` diff --git a/docs/developer/architecture.rst b/docs/developer/architecture.rst index ccbf06b..a8b2f55 100644 --- a/docs/developer/architecture.rst +++ b/docs/developer/architecture.rst @@ -38,8 +38,8 @@ Workflow Classes Orchestrates 4D CT loading, segmentation, image registration, contour transformation, and animated USD generation. -``WorkflowConvertCTToVTK`` - Converts one CT image into labeled VTK surface and voxel-mesh outputs. +``WorkflowConvertImageToVTK`` + Converts one 3D image into labeled VTK surface and voxel-mesh outputs. ``WorkflowCreateStatisticalModel`` Builds a PCA statistical shape model from aligned population meshes. diff --git a/docs/developer/workflows.rst b/docs/developer/workflows.rst index a8869d5..5c0f8e8 100644 --- a/docs/developer/workflows.rst +++ b/docs/developer/workflows.rst @@ -16,8 +16,8 @@ Current Workflow Mapping - Workflow class * - ``physiomotion4d-convert-image-to-usd`` - :class:`physiomotion4d.WorkflowConvertImageToUSD` - * - ``physiomotion4d-convert-ct-to-vtk`` - - :class:`physiomotion4d.WorkflowConvertCTToVTK` + * - ``physiomotion4d-convert-image-to-vtk`` + - :class:`physiomotion4d.WorkflowConvertImageToVTK` * - ``physiomotion4d-convert-vtk-to-usd`` - :class:`physiomotion4d.WorkflowConvertVTKToUSD` * - ``physiomotion4d-create-statistical-model`` @@ -39,7 +39,7 @@ Workflow Example contrast_enhanced=True, output_directory="./results", project_name="patient_001", - registration_method="ants", + registration_method="ANTS", ) final_usd = workflow.process() diff --git a/docs/examples.rst b/docs/examples.rst index 3ca3d99..283e018 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -35,7 +35,7 @@ Complete end-to-end cardiac CT processing: contrast_enhanced=True, output_directory="./results", project_name="patient_001", - registration_method="ants", + registration_method="ANTS", ) # Run complete workflow @@ -411,7 +411,7 @@ Batch process multiple datasets: contrast_enhanced=True, output_directory=f"results/{patient_id}", project_name=patient_id, - registration_method="ants", + registration_method="ANTS", ) try: @@ -481,7 +481,7 @@ Run the supported end-to-end workflow API: contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", - registration_method="ants", + registration_method="ANTS", ) final_usd = workflow.process() diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 68fed1c..d719f87 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -71,7 +71,7 @@ CUDA-capable GPU are required for practical runtime. python -c "from physiomotion4d import DataDownloadTools; DataDownloadTools.DownloadSlicerHeartCTData('data/test')" physiomotion4d-convert-image-to-usd data/test/TruncalValve_4DCT.seq.nrrd \ - --registration-method ants \ + --registration-method ANTS \ --output-dir output/quickstart \ --project-name slicer_heart_quickstart @@ -115,7 +115,7 @@ For more control, use the Python API: contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", - registration_method="ants", + registration_method="ANTS", ) **Step 3: Run the workflow** @@ -150,7 +150,7 @@ For more control over individual steps: contrast_enhanced=True, output_directory="./results", project_name="cardiac_model", - registration_method="ants", + registration_method="ANTS", ) final_usd = workflow.process() @@ -312,7 +312,7 @@ Common Issues * Resample or crop the input image before running the workflow * Process fewer frames at once -* Use ANTs registration with ``--registration-method ants`` when CUDA is unavailable +* Use ANTs registration with ``--registration-method ANTS`` when CUDA is unavailable **Segmentation quality issues** diff --git a/docs/troubleshooting.rst b/docs/troubleshooting.rst index 659d733..546b954 100644 --- a/docs/troubleshooting.rst +++ b/docs/troubleshooting.rst @@ -15,7 +15,7 @@ CUDA Out of Memory **Solutions**: 1. Resample or crop the input image before running the workflow. -2. Use ``--registration-method ants`` when CUDA is unavailable. +2. Use ``--registration-method ANTS`` when CUDA is unavailable. 3. Process fewer frames per run. CUDA Version Mismatch @@ -93,7 +93,7 @@ Registration Not Converging .. code-block:: bash - physiomotion4d-convert-image-to-usd cardiac_4d.nrrd --registration-method ants + physiomotion4d-convert-image-to-usd cardiac_4d.nrrd --registration-method ANTS 3. Check image orientation and spacing diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 781884e..0634ee8 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -121,7 +121,7 @@ Script ``tutorials/tutorial_02_ct_to_vtk.py`` Workflow - ``WorkflowConvertCTToVTK`` + ``WorkflowConvertImageToVTK`` Dataset Slicer-Heart-CT, prepared before running the tutorial. @@ -222,7 +222,7 @@ Script ``tutorials/tutorial_07_dirlab_pca_model.py`` Workflow - ``WorkflowConvertCTToVTK``, ``WorkflowCreateStatisticalModel``, and + ``WorkflowConvertImageToVTK``, ``WorkflowCreateStatisticalModel``, and ``WorkflowFitStatisticalModelToPatient`` Dataset diff --git a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py index 4f40d74..0def33e 100644 --- a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py +++ b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py @@ -274,7 +274,7 @@ heart_array = itk.array_from_image(result["heart"]) vessels_array = itk.array_from_image(result["major_vessels"]) - labelmap_essentials = segmenter.trim_mask_to_essentials(result["labelmap"]) + labelmap_essentials = segmenter.trim_branches(result["labelmap"]) labelmap_essentials_array = itk.array_from_image(labelmap_essentials) # Select middle slice diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py index 0c00e15..3a2c53b 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py @@ -51,7 +51,7 @@ registrar = WorkflowFitStatisticalModelToPatient( template_model=template_model, patient_image=patient_image, - segmentation_method="simpleware_heart", + segmentation_method="HeartSimpleware", ) registrar.set_use_pca_registration( diff --git a/pyproject.toml b/pyproject.toml index 379c8f0..11c2f61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,7 +160,7 @@ Repository = "https://github.com/Project-MONAI/physiomotion4d.git" [project.scripts] # CLI commands installed via pip # Entry points reference the main() functions in the cli submodule -physiomotion4d-convert-ct-to-vtk = "physiomotion4d.cli.convert_ct_to_vtk:main" +physiomotion4d-convert-image-to-vtk = "physiomotion4d.cli.convert_image_to_vtk:main" physiomotion4d-convert-image-4d-to-3d = "physiomotion4d.cli.convert_image_4d_to_3d:main" physiomotion4d-convert-image-to-usd = "physiomotion4d.cli.convert_image_to_usd:main" physiomotion4d-convert-vtk-to-usd = "physiomotion4d.cli.convert_vtk_to_usd:main" @@ -236,6 +236,8 @@ module = [ "SimpleITK", "SimpleITK.*", "sklearn.*", + "torch", + "torch.*", "totalsegmentator.*", "trimesh", "trimesh.*", diff --git a/src/physiomotion4d/__init__.py b/src/physiomotion4d/__init__.py index c1f5071..9fb2c24 100644 --- a/src/physiomotion4d/__init__.py +++ b/src/physiomotion4d/__init__.py @@ -68,7 +68,7 @@ from .usd_tools import USDTools # Core workflow processor -from .workflow_convert_ct_to_vtk import WorkflowConvertCTToVTK +from .workflow_convert_image_to_vtk import WorkflowConvertImageToVTK from .workflow_convert_image_to_usd import WorkflowConvertImageToUSD from .workflow_convert_vtk_to_usd import WorkflowConvertVTKToUSD from .workflow_reconstruct_highres_4d_ct import WorkflowReconstructHighres4DCT @@ -79,7 +79,7 @@ __all__ = [ # Workflow classes - "WorkflowConvertCTToVTK", + "WorkflowConvertImageToVTK", "WorkflowConvertImageToUSD", "WorkflowConvertVTKToUSD", "WorkflowCreateStatisticalModel", diff --git a/src/physiomotion4d/cli/convert_image_to_usd.py b/src/physiomotion4d/cli/convert_image_to_usd.py index 317b130..58c5f0c 100644 --- a/src/physiomotion4d/cli/convert_image_to_usd.py +++ b/src/physiomotion4d/cli/convert_image_to_usd.py @@ -28,7 +28,10 @@ def main() -> int: %(prog)s input.nrrd --reference-image ref.mha --registration-iterations 50 # Use ANTs registration instead of ICON - %(prog)s input.nrrd --contrast --registration-method ants + %(prog)s input.nrrd --contrast --registration-method ANTS + + # Use the cardiac-only Simpleware segmentation backend + %(prog)s input.nrrd --segmentation-method HeartSimpleware """, ) @@ -64,11 +67,20 @@ def main() -> int: default=1, help="Number of registration iterations (default: 1)", ) + parser.add_argument( + "--segmentation-method", + choices=["ChestTotalSegmentator", "HeartSimpleware"], + default="ChestTotalSegmentator", + help=( + "Segmentation backend to use: ChestTotalSegmentator (default) " + "or HeartSimpleware." + ), + ) parser.add_argument( "--registration-method", - choices=["ants", "icon"], - default="icon", - help="Registration method to use: ants or icon (default: icon)", + choices=["ANTS", "ICON"], + default="ICON", + help="Registration method to use: ANTS or ICON (default: ICON)", ) args = parser.parse_args() @@ -91,6 +103,7 @@ def main() -> int: project_name=args.project_name, reference_image_filename=args.reference_image, number_of_registration_iterations=args.registration_iterations, + segmentation_method=args.segmentation_method, registration_method=args.registration_method, ) except Exception as e: diff --git a/src/physiomotion4d/cli/convert_ct_to_vtk.py b/src/physiomotion4d/cli/convert_image_to_vtk.py similarity index 88% rename from src/physiomotion4d/cli/convert_ct_to_vtk.py rename to src/physiomotion4d/cli/convert_image_to_vtk.py index b4d91b3..ff29872 100644 --- a/src/physiomotion4d/cli/convert_ct_to_vtk.py +++ b/src/physiomotion4d/cli/convert_image_to_vtk.py @@ -1,7 +1,7 @@ #!/usr/bin/env python -"""Command-line interface for the CT-to-VTK segmentation workflow. +"""Command-line interface for the image-to-VTK segmentation workflow. -Segments a 3D CT image using a chosen backend and writes per-anatomy-group VTP +Segments a 3D image using a chosen backend and writes per-anatomy-group VTP surfaces and VTU voxel meshes annotated with anatomy labels and colors. """ @@ -20,13 +20,13 @@ "contrast", ) -SEGMENTATION_METHODS = ("total_segmentator", "simpleware_heart") +SEGMENTATION_METHODS = ("ChestTotalSegmentator", "HeartSimpleware") def main() -> int: - """CLI entry point for CT to VTK conversion.""" + """CLI entry point for image to VTK conversion.""" parser = argparse.ArgumentParser( - description="Segment a CT image and export anatomy groups as VTK surfaces and meshes.", + description="Segment a 3D image and export anatomy groups as VTK surfaces and meshes.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Anatomy groups @@ -54,7 +54,7 @@ def main() -> int: # Simpleware heart-only, cardiac anatomy groups, combined output %(prog)s \\ --input-image chest_ct.nii.gz \\ - --segmentation-method simpleware_heart \\ + --segmentation-method HeartSimpleware \\ --anatomy-groups heart major_vessels \\ --output-dir ./results \\ --output-prefix patient01 @@ -71,7 +71,7 @@ def main() -> int: parser.add_argument( "--input-image", required=True, - help="Path to the input CT image (.nii.gz, .nrrd, .mha, …).", + help="Path to the input 3D image (.nii.gz, .nrrd, .mha, …).", ) parser.add_argument( "--output-dir", @@ -82,9 +82,11 @@ def main() -> int: # ── Segmentation ────────────────────────────────────────────────────── parser.add_argument( "--segmentation-method", - default="total_segmentator", + default="ChestTotalSegmentator", choices=list(SEGMENTATION_METHODS), - help=("Segmentation backend. total_segmentator (default) | simpleware_heart"), + help=( + "Segmentation backend. ChestTotalSegmentator (default) | HeartSimpleware" + ), ) parser.add_argument( "--contrast", @@ -151,9 +153,9 @@ def main() -> int: print("=" * 70) try: - from physiomotion4d import WorkflowConvertCTToVTK + from physiomotion4d import WorkflowConvertImageToVTK - workflow = WorkflowConvertCTToVTK( + workflow = WorkflowConvertImageToVTK( segmentation_method=args.segmentation_method, ) result = workflow.run_workflow( @@ -183,13 +185,13 @@ def main() -> int: if args.split_files: # One file per anatomy group if surfaces: - saved_surfaces = WorkflowConvertCTToVTK.save_surfaces( + saved_surfaces = WorkflowConvertImageToVTK.save_surfaces( surfaces, args.output_dir, prefix=prefix ) for group, path in saved_surfaces.items(): print(f" Surface [{group:15s}] → {path}") if meshes: - saved_meshes = WorkflowConvertCTToVTK.save_meshes( + saved_meshes = WorkflowConvertImageToVTK.save_meshes( meshes, args.output_dir, prefix=prefix ) for group, path in saved_meshes.items(): @@ -197,12 +199,12 @@ def main() -> int: else: # Combined single-file output if surfaces: - surface_file = WorkflowConvertCTToVTK.save_combined_surface( + surface_file = WorkflowConvertImageToVTK.save_combined_surface( surfaces, args.output_dir, prefix=prefix ) print(f" Combined surface → {surface_file}") if meshes: - mesh_file = WorkflowConvertCTToVTK.save_combined_mesh( + mesh_file = WorkflowConvertImageToVTK.save_combined_mesh( meshes, args.output_dir, prefix=prefix ) print(f" Combined mesh → {mesh_file}") diff --git a/src/physiomotion4d/register_images_icon.py b/src/physiomotion4d/register_images_icon.py index 9cc93a7..31072cf 100644 --- a/src/physiomotion4d/register_images_icon.py +++ b/src/physiomotion4d/register_images_icon.py @@ -10,17 +10,24 @@ """ import logging -from typing import Optional, Union +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Optional, Union import icon_registration as icon import icon_registration.itk_wrapper import itk +import numpy as np +import torch +import torch.nn.functional as F from unigradicon import get_multigradicon, get_unigradicon from unigradicon import preprocess as unigradicon_preprocess from physiomotion4d.register_images_base import RegisterImagesBase from physiomotion4d.transform_tools import TransformTools +DEFAULT_FINETUNE_LEARNING_RATE = 2e-5 + class RegisterImagesICON(RegisterImagesBase): """ICON-based deformable image registration implementation. @@ -226,20 +233,7 @@ def registration_method( moving_effective_mask = moving_labelmap if use_labelmaps else moving_mask fixed_effective_mask = self.fixed_labelmap if use_labelmaps else self.fixed_mask - if self.net is None: - if self.use_multi_modality: - self.net = get_multigradicon( - loss_fn=icon.LNCC(sigma=5), - # loss_fn=icon.losses.MINDSSC(radius=2, dilation=2), - apply_intensity_conservation_loss=self.use_mass_preservation, - weights_location=self.weights_path, - ) - else: - self.net = get_unigradicon( - loss_fn=icon.LNCC(sigma=5), - apply_intensity_conservation_loss=self.use_mass_preservation, - weights_location=self.weights_path, - ) + self._ensure_net() inverse_transform = None forward_transform = None @@ -290,3 +284,159 @@ def registration_method( "inverse_transform": inverse_transform, "loss": loss, } + + def _ensure_net(self) -> None: + """Lazily instantiate the ICON network using current configuration. + + Honors set_weights_path() if a custom checkpoint has been requested, + otherwise loads the default UniGradICON / MultiGradICON pretrained + weights. + """ + if self.net is not None: + return + if self.use_multi_modality: + self.net = get_multigradicon( + loss_fn=icon.LNCC(sigma=5), + apply_intensity_conservation_loss=self.use_mass_preservation, + weights_location=self.weights_path, + ) + else: + self.net = get_unigradicon( + loss_fn=icon.LNCC(sigma=5), + apply_intensity_conservation_loss=self.use_mass_preservation, + weights_location=self.weights_path, + ) + + def _image_to_resized_tensor( + self, image: itk.Image, shape: torch.Size + ) -> torch.Tensor: + """Convert an itk image to a torch tensor resized to the net's input grid. + + Mirrors the trilinear preprocessing path used by + ``icon_registration.itk_wrapper.register_pair``. + """ + arr = np.array(image) + tensor = torch.Tensor(arr).to(icon.config.device)[None, None] + return F.interpolate( + tensor, size=shape[2:], mode="trilinear", align_corners=False + ) + + def _mask_to_resized_tensor( + self, mask: itk.Image, shape: torch.Size + ) -> torch.Tensor: + """Convert an itk mask image to a torch tensor resized via nearest-neighbor. + + Mirrors the mask preprocessing used by + ``icon_registration.itk_wrapper.register_pair_with_mask``. + """ + arr = np.array(mask) + tensor = torch.Tensor(arr).to(icon.config.device)[None, None] + return F.interpolate(tensor, size=shape[2:], mode="nearest") + + def finetune( + self, + image_pairs: Sequence[tuple[itk.Image, itk.Image]], + output_model_filename: str, + mask_pairs: Optional[Sequence[tuple[itk.Image, itk.Image]]] = None, + epochs: int = 1, + learning_rate: float = DEFAULT_FINETUNE_LEARNING_RATE, + ) -> dict[str, Any]: + """Fine-tune the ICON network on a cohort of image pairs. + + Unlike ``register()``, this method *persistently* updates the in-memory + network weights and saves the resulting inner-network state_dict to + disk. The starting weights are whatever ``set_weights_path()`` was + configured to (default UniGradICON pretrained weights if never set). + + This method intentionally does not call ``register()`` / the ICON + ``register_pair`` helpers, because those wrap their optimization in a + ``state_dict`` save/restore that discards any persistent weight + changes. Here the optimizer steps act directly on ``self.net`` and + are not undone. + + Args: + image_pairs: Sequence of (fixed_image, moving_image) itk image + pairs to fine-tune on. Caller can build this from a list of + images with ``itertools.combinations(images, 2)``. + output_model_filename: Path where the fine-tuned inner-network + state_dict will be saved. Suitable for reloading via + ``set_weights_path()`` on a new RegisterImagesICON instance. + mask_pairs: Optional sequence of (fixed_mask, moving_mask) itk + image pairs, same length as ``image_pairs``. When provided, + ICON's masked-loss path is used for every pair. + epochs: Number of full passes over ``image_pairs``. Default 1. + learning_rate: Adam learning rate. Default matches + ``icon_registration``'s per-pair fine-tune rate. + + Returns: + Dict with keys: + - ``output_model_filename``: str path of saved checkpoint + - ``losses``: list[float], one entry per (epoch, pair) + - ``epochs``: int + - ``number_of_pairs``: int + + Raises: + ValueError: If ``image_pairs`` is empty or ``mask_pairs`` length + does not match ``image_pairs``. + """ + if len(image_pairs) == 0: + raise ValueError("At least one image pair is required for finetune.") + if mask_pairs is not None and len(mask_pairs) != len(image_pairs): + raise ValueError("mask_pairs must match image_pairs length.") + if epochs < 1: + raise ValueError("epochs must be >= 1.") + + self._ensure_net() + assert self.net is not None + self.net.to(icon.config.device) + self.net.train() + + optimizer = torch.optim.Adam(self.net.parameters(), lr=learning_rate) + shape = self.net.identity_map.shape + losses: list[float] = [] + + for epoch_index in range(epochs): + for pair_index, (fixed_image, moving_image) in enumerate(image_pairs): + fixed_pre = self.preprocess(fixed_image, modality=self.modality) + moving_pre = self.preprocess(moving_image, modality=self.modality) + + fixed_resized = self._image_to_resized_tensor(fixed_pre, shape) + moving_resized = self._image_to_resized_tensor(moving_pre, shape) + + forward_kwargs: dict[str, torch.Tensor] = {} + if mask_pairs is not None: + fixed_mask, moving_mask = mask_pairs[pair_index] + forward_kwargs["mask_A"] = self._mask_to_resized_tensor( + fixed_mask, shape + ) + forward_kwargs["mask_B"] = self._mask_to_resized_tensor( + moving_mask, shape + ) + + optimizer.zero_grad() + loss_tuple = self.net(fixed_resized, moving_resized, **forward_kwargs) + loss = loss_tuple[0] + loss.backward() + optimizer.step() + + loss_value = float(loss.detach().cpu().item()) + losses.append(loss_value) + self.log_info( + "finetune epoch %d pair %d/%d loss=%.6f", + epoch_index, + pair_index + 1, + len(image_pairs), + loss_value, + ) + + output_path = Path(output_model_filename) + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(self.net.regis_net.state_dict(), output_path) + self.weights_path = str(output_path) + + return { + "output_model_filename": str(output_path), + "losses": losses, + "epochs": epochs, + "number_of_pairs": len(image_pairs), + } diff --git a/src/physiomotion4d/segment_heart_simpleware.py b/src/physiomotion4d/segment_heart_simpleware.py index 2b39a18..c9eec86 100644 --- a/src/physiomotion4d/segment_heart_simpleware.py +++ b/src/physiomotion4d/segment_heart_simpleware.py @@ -115,13 +115,20 @@ def __init__(self, log_level: int | str = logging.INFO): "SimplewareScript_heart_segmentation.py", ) - def set_trim_mask_to_essentials(self, trim_mask: bool) -> None: - """Set whether to trim mask to common and critical structures. + def set_trim_branches(self, trim_branches: bool) -> None: + """Enable trimming of pulmonary and great-vessel branches. + + When enabled, :meth:`trim_branches` is applied to the labelmap after + Simpleware segmentation, clipping the pulmonary veins and major great + vessels back to the cardiac region. Trimming reduces inter-subject + variability in vessel extent, which simplifies AI-Ready and Sim-Ready + model fitting. It is consistent with how vessels were trimmed in the + example KCL Heart dataset. Args: - trim_mask (bool): Whether to reduce to essential. + trim_branches (bool): Whether to trim branches to the cardiac region. """ - self._trim_mask = trim_mask + self._trim_mask = trim_branches def set_simpleware_executable_path(self, path: str) -> None: """Set the path to the Simpleware Medical console executable. @@ -332,7 +339,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: labelmap_image.CopyInformation(preprocessed_image) if self._trim_mask: - labelmap_image = self.trim_mask_to_essentials(labelmap_image) + labelmap_image = self.trim_branches(labelmap_image) return labelmap_image @@ -340,8 +347,16 @@ def get_landmarks(self) -> dict[str, tuple[float, float, float]]: """Get the landmarks.""" return self.landmarks - def trim_mask_to_essentials(self, labelmap_image: itk.image) -> itk.image: - """Trim mask to essentials.""" + def trim_branches(self, labelmap_image: itk.image) -> itk.image: + """Trim pulmonary and great-vessel branches back to the cardiac region. + + Clips pulmonary veins and the great vessels (aorta, pulmonary artery) + to the portions adjacent to the heart and keeps only the largest + connected component of the left and right atria. Reduces inter-subject + variability in vessel extent, simplifying AI-Ready and Sim-Ready model + fitting. Consistent with how vessels were trimmed in the example KCL + Heart dataset. + """ # Reference code for cropping aorta and pulmonary artery to # portions adjacent to the heart. diff --git a/src/physiomotion4d/workflow_convert_image_to_usd.py b/src/physiomotion4d/workflow_convert_image_to_usd.py index eb3f853..e19dfce 100644 --- a/src/physiomotion4d/workflow_convert_image_to_usd.py +++ b/src/physiomotion4d/workflow_convert_image_to_usd.py @@ -22,10 +22,21 @@ from physiomotion4d.register_images_ants import RegisterImagesANTs from physiomotion4d.register_images_base import RegisterImagesBase from physiomotion4d.register_images_icon import RegisterImagesICON +from physiomotion4d.segment_anatomy_base import SegmentAnatomyBase from physiomotion4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator +from physiomotion4d.segment_heart_simpleware import SegmentHeartSimpleware from physiomotion4d.transform_tools import TransformTools from physiomotion4d.usd_anatomy_tools import USDAnatomyTools +#: Supported segmentation backend identifiers. +SEGMENTATION_METHODS: tuple[str, ...] = ( + "ChestTotalSegmentator", + "HeartSimpleware", +) + +#: Supported registration backend identifiers. +REGISTRATION_METHODS: tuple[str, ...] = ("ANTS", "ICON") + class WorkflowConvertImageToUSD(PhysioMotion4DBase): """ @@ -43,7 +54,8 @@ def __init__( project_name: str, reference_image_filename: Optional[str] = None, number_of_registration_iterations: Optional[int] = 1, - registration_method: str = "icon", + segmentation_method: str = "ChestTotalSegmentator", + registration_method: str = "ICON", log_level: int | str = logging.INFO, save_registered_images: bool = True, save_registration_transforms: bool = True, @@ -65,7 +77,10 @@ def __init__( project_name (str): Project name for USD file organization reference_image_filename (Optional[str]): Path to reference image file number_of_registration_iterations (Optional[int]): Number of registration iterations - registration_method (str): Registration method to use: 'ants' or 'icon' (default: 'icon') + segmentation_method (str): Segmentation backend to use: + ``'ChestTotalSegmentator'`` (default) or ``'HeartSimpleware'``. + registration_method (str): Registration method to use: + ``'ANTS'`` or ``'ICON'`` (default: ``'ICON'``). log_level: Logging level (default: logging.INFO) save_registered_images: Write registered image intermediates to output_directory when True @@ -86,11 +101,19 @@ def __init__( self.save_registration_transforms = save_registration_transforms self.save_labelmaps = save_labelmaps + # Validate segmentation method + if segmentation_method not in SEGMENTATION_METHODS: + raise ValueError( + f"Invalid segmentation_method '{segmentation_method}'. " + f"Must be one of: {', '.join(SEGMENTATION_METHODS)}." + ) + self.segmentation_method = segmentation_method + # Validate registration method - if registration_method not in ["ants", "icon"]: + if registration_method not in REGISTRATION_METHODS: raise ValueError( f"Invalid registration_method '{registration_method}'. " - "Must be 'ants' or 'icon'." + f"Must be one of: {', '.join(REGISTRATION_METHODS)}." ) self.registration_method = registration_method @@ -99,12 +122,19 @@ def __init__( # Initialize processing components self.converter = ConvertImage4DTo3D(log_level=log_level) - self.segmenter = SegmentChestTotalSegmentator(log_level=log_level) - self.segmenter.contrast_threshold = 500 + self.segmenter: SegmentAnatomyBase + if self.segmentation_method == "ChestTotalSegmentator": + chest_segmenter = SegmentChestTotalSegmentator(log_level=log_level) + chest_segmenter.contrast_threshold = 500 + self.segmenter = chest_segmenter + else: # HeartSimpleware + heart_segmenter = SegmentHeartSimpleware(log_level=log_level) + heart_segmenter.set_trim_branches(True) + self.segmenter = heart_segmenter # Initialize registration method self.registrar: RegisterImagesBase - if self.registration_method == "ants": + if self.registration_method == "ANTS": self.log_info("Initializing ANTs registration...") ants_registrar = RegisterImagesANTs(log_level=log_level) ants_registrar.set_modality("ct") @@ -121,7 +151,7 @@ def __init__( ] ) self.registrar = ants_registrar - else: # icon (default) + else: # ICON (default) self.log_info("Initializing ICON registration...") icon_registrar = RegisterImagesICON(log_level=log_level) icon_registrar.set_modality("ct") @@ -294,7 +324,7 @@ def _segment_and_register_frames(self) -> None: # Set up registrar with fixed image self.registrar.set_fixed_image(self._fixed_image) - if self.registration_method == "icon" and isinstance( + if self.registration_method == "ICON" and isinstance( self.registrar, RegisterImagesICON ): if self.contrast_enhanced: diff --git a/src/physiomotion4d/workflow_convert_ct_to_vtk.py b/src/physiomotion4d/workflow_convert_image_to_vtk.py similarity index 92% rename from src/physiomotion4d/workflow_convert_ct_to_vtk.py rename to src/physiomotion4d/workflow_convert_image_to_vtk.py index 5283ff0..2991ddf 100644 --- a/src/physiomotion4d/workflow_convert_ct_to_vtk.py +++ b/src/physiomotion4d/workflow_convert_image_to_vtk.py @@ -9,19 +9,19 @@ Typical usage:: import itk - from physiomotion4d import WorkflowConvertCTToVTK + from physiomotion4d import WorkflowConvertImageToVTK ct = itk.imread('chest_ct.nii.gz') - workflow = WorkflowConvertCTToVTK(segmentation_method='total_segmentator') + workflow = WorkflowConvertImageToVTK(segmentation_method='ChestTotalSegmentator') result = workflow.run_workflow(ct, contrast_enhanced_study=True) # Combined single-file output (default) - WorkflowConvertCTToVTK.save_combined_surface(result['surfaces'], './out', prefix='patient') - WorkflowConvertCTToVTK.save_combined_mesh(result['meshes'], './out', prefix='patient') + WorkflowConvertImageToVTK.save_combined_surface(result['surfaces'], './out', prefix='patient') + WorkflowConvertImageToVTK.save_combined_mesh(result['meshes'], './out', prefix='patient') # Per-group split output - WorkflowConvertCTToVTK.save_surfaces(result['surfaces'], './out', prefix='patient') - WorkflowConvertCTToVTK.save_meshes(result['meshes'], './out', prefix='patient') + WorkflowConvertImageToVTK.save_surfaces(result['surfaces'], './out', prefix='patient') + WorkflowConvertImageToVTK.save_meshes(result['meshes'], './out', prefix='patient') """ import logging @@ -50,20 +50,20 @@ #: Supported segmentation backend identifiers. SEGMENTATION_METHODS: tuple[str, ...] = ( - "total_segmentator", - "simpleware_heart", + "ChestTotalSegmentator", + "HeartSimpleware", ) -class WorkflowConvertCTToVTK(PhysioMotion4DBase): +class WorkflowConvertImageToVTK(PhysioMotion4DBase): """Segment a CT image and produce per-anatomy-group VTK surfaces and meshes. **Segmentation backends** - - ``'total_segmentator'`` — :class:`SegmentChestTotalSegmentator` (CPU-capable, - default). - - ``'simpleware_heart'`` — :class:`SegmentHeartSimpleware` (cardiac only; requires - a Simpleware Medical installation). + - ``'ChestTotalSegmentator'`` — :class:`SegmentChestTotalSegmentator` + (CPU-capable, default). + - ``'HeartSimpleware'`` — :class:`SegmentHeartSimpleware` (cardiac only; + requires a Simpleware Medical installation). **Output anatomy groups** @@ -86,7 +86,7 @@ class WorkflowConvertCTToVTK(PhysioMotion4DBase): :meth:`run_workflow` performs *no* file I/O. Use the static helpers :meth:`save_surfaces`, :meth:`save_meshes`, :meth:`save_combined_surface`, and - :meth:`save_combined_mesh` — or the CLI ``physiomotion4d-convert-ct-to-vtk`` — to + :meth:`save_combined_mesh` — or the CLI ``physiomotion4d-convert-image-to-vtk`` — to write results to disk. """ @@ -97,14 +97,14 @@ class WorkflowConvertCTToVTK(PhysioMotion4DBase): def __init__( self, - segmentation_method: str = "total_segmentator", + segmentation_method: str = "ChestTotalSegmentator", log_level: int | str = logging.INFO, ) -> None: """Initialize the workflow. Args: segmentation_method: Segmentation backend to use. One of - ``'total_segmentator'`` (default) or ``'simpleware_heart'``. + ``'ChestTotalSegmentator'`` (default) or ``'HeartSimpleware'``. log_level: Logging level. Default: ``logging.INFO``. Raises: @@ -138,17 +138,17 @@ def __init__( def _create_segmenter(self) -> SegmentAnatomyBase: """Instantiate the chosen segmentation backend (lazy import).""" - if self.segmentation_method_name == "total_segmentator": + if self.segmentation_method_name == "ChestTotalSegmentator": from physiomotion4d.segment_chest_total_segmentator import ( SegmentChestTotalSegmentator, ) return SegmentChestTotalSegmentator(log_level=self.log_level) - if self.segmentation_method_name == "simpleware_heart": + if self.segmentation_method_name == "HeartSimpleware": from physiomotion4d.segment_heart_simpleware import SegmentHeartSimpleware segmenter = SegmentHeartSimpleware(log_level=self.log_level) - segmenter.set_trim_mask_to_essentials(True) + segmenter.set_trim_branches(True) return segmenter raise ValueError( f"Unknown segmentation method: {self.segmentation_method_name}" @@ -158,7 +158,7 @@ def _get_label_info_for_group(self, group: str) -> tuple[list[str], list[int]]: """Return ``(label_names, label_ids)`` for *group* from the active segmenter. Reads the segmenter's :class:`AnatomyTaxonomy`. Returns empty lists if - the group is not present (e.g. simpleware_heart does not register + the group is not present (e.g. HeartSimpleware does not register lung/bone). """ assert self._segmenter is not None, ( @@ -268,7 +268,7 @@ def run_workflow( Raises: ValueError: If any name in *anatomy_groups* is invalid. """ - self.log_section("STARTING CT TO VTK WORKFLOW") + self.log_section("STARTING IMAGE TO VTK WORKFLOW") # Validate requested groups if anatomy_groups is not None: @@ -327,7 +327,7 @@ def run_workflow( self._annotate(mesh, group, label_names, label_ids, color) meshes[group] = mesh - self.log_section("CT TO VTK WORKFLOW COMPLETE") + self.log_section("IMAGE TO VTK WORKFLOW COMPLETE") self.log_info("Surfaces extracted: %d", len(surfaces)) self.log_info("Meshes extracted: %d", len(meshes)) diff --git a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py index 77f8c16..5cacd62 100644 --- a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py @@ -37,7 +37,7 @@ from physiomotion4d.register_models_icp import RegisterModelsICP from physiomotion4d.register_models_pca import RegisterModelsPCA from physiomotion4d.transform_tools import TransformTools -from physiomotion4d.workflow_convert_ct_to_vtk import WorkflowConvertCTToVTK +from physiomotion4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK def _load_tubetk() -> Any: @@ -137,7 +137,7 @@ def __init__( template_model: pv.DataSet, patient_models: list[pv.DataSet] | None = None, patient_image: Optional[itk.Image] = None, - segmentation_method: str = "simpleware_heart", + segmentation_method: str = "HeartSimpleware", log_level: int | str = logging.INFO, ): """Initialize the model-to-image-and-model registration pipeline. @@ -167,11 +167,11 @@ def __init__( self.template_labelmap_background_ids: Optional[list[int]] = None if patient_models is None and patient_image is not None: - convert_ct_to_vtk = WorkflowConvertCTToVTK( + convert_image_to_vtk = WorkflowConvertImageToVTK( segmentation_method=segmentation_method, log_level=log_level, ) - patient_models_data = convert_ct_to_vtk.run_workflow( + patient_models_data = convert_image_to_vtk.run_workflow( input_image=patient_image, contrast_enhanced_study=False, anatomy_groups=["heart"], diff --git a/tests/test_cli_smoke.py b/tests/test_cli_smoke.py index b130298..dcb236a 100644 --- a/tests/test_cli_smoke.py +++ b/tests/test_cli_smoke.py @@ -9,7 +9,7 @@ CLI_MODULES = [ - "physiomotion4d.cli.convert_ct_to_vtk", + "physiomotion4d.cli.convert_image_to_vtk", "physiomotion4d.cli.convert_image_4d_to_3d", "physiomotion4d.cli.convert_image_to_usd", "physiomotion4d.cli.convert_vtk_to_usd", diff --git a/tutorials/README.md b/tutorials/README.md index bd24303..805988d 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -14,7 +14,7 @@ dataset licensing, and expected directory layout. | # | Script | Primary API | Dataset | |---|--------|-------------|---------| | 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | -| 2 | [tutorial_02_ct_to_vtk.py](tutorial_02_ct_to_vtk.py) | `WorkflowConvertCTToVTK` | Slicer-Heart-CT (prepare first) | +| 2 | [tutorial_02_ct_to_vtk.py](tutorial_02_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | | 3 | [tutorial_03_create_statistical_model.py](tutorial_03_create_statistical_model.py) | `WorkflowCreateStatisticalModel` | KCL-Heart-Model (manual) | | 4 | [tutorial_04_fit_statistical_model_to_patient.py](tutorial_04_fit_statistical_model_to_patient.py) | `WorkflowFitStatisticalModelToPatient` | KCL-Heart-Model plus Tutorial 3 output | | 5 | [tutorial_05_vtk_to_usd.py](tutorial_05_vtk_to_usd.py) | `WorkflowConvertVTKToUSD` | Output of tutorial 2 | diff --git a/tutorials/tutorial_01_heart_gated_ct_to_usd.py b/tutorials/tutorial_01_heart_gated_ct_to_usd.py index b732789..19281cb 100644 --- a/tutorials/tutorial_01_heart_gated_ct_to_usd.py +++ b/tutorials/tutorial_01_heart_gated_ct_to_usd.py @@ -36,8 +36,8 @@ Weaknesses / Limitations ------------------------ -- Requires a GPU for ICON registration (``registration_method='icon'``); use - ``registration_method='ants'`` for CPU-only environments (about 10x slower). +- Requires a GPU for ICON registration (``registration_method='ICON'``); use + ``registration_method='ANTS'`` for CPU-only environments (about 10x slower). - Segmentation quality depends on TotalSegmentator's training distribution; unusual pathologies or pediatric anatomy may degrade results. - Large 4D datasets (>20 phases, high resolution) can require 32 GB+ RAM. @@ -93,7 +93,7 @@ FULL_DATA_DIR = DATA_DIR / "Slicer-Heart-CT" TEST_DATA_DIR = DATA_DIR / "test" / "slicer_heart_small" OUTPUT_DIR = TUTORIALS_DIR / "output" / "tutorial_01" - REGISTRATION_METHOD = "ants" + REGISTRATION_METHOD = "ANTS" LOG_LEVEL = logging.INFO # %% diff --git a/tutorials/tutorial_02_ct_to_vtk.py b/tutorials/tutorial_02_ct_to_vtk.py index 02c406a..a961a36 100644 --- a/tutorials/tutorial_02_ct_to_vtk.py +++ b/tutorials/tutorial_02_ct_to_vtk.py @@ -24,7 +24,7 @@ import pyvista as pv from physiomotion4d.test_tools import TestTools -from physiomotion4d.workflow_convert_ct_to_vtk import WorkflowConvertCTToVTK +from physiomotion4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK # nnUNetv2 (used by TotalSegmentator inside several workflows) spawns a # multiprocessing.Pool. On Windows the spawn start method re-imports this @@ -67,8 +67,8 @@ # %% # Workflow initialization - workflow = WorkflowConvertCTToVTK( - segmentation_method="total_segmentator", + workflow = WorkflowConvertImageToVTK( + segmentation_method="ChestTotalSegmentator", log_level=log_level, ) @@ -82,14 +82,14 @@ # %% # Result saving surface_file = Path( - WorkflowConvertCTToVTK.save_combined_surface( + WorkflowConvertImageToVTK.save_combined_surface( result["surfaces"], str(output_dir), prefix="patient", ) ) mesh_file = Path( - WorkflowConvertCTToVTK.save_combined_mesh( + WorkflowConvertImageToVTK.save_combined_mesh( result["meshes"], str(output_dir), prefix="patient", diff --git a/tutorials/tutorial_07_dirlab_pca_model.py b/tutorials/tutorial_07_dirlab_pca_model.py index edfae3d..b0dc589 100644 --- a/tutorials/tutorial_07_dirlab_pca_model.py +++ b/tutorials/tutorial_07_dirlab_pca_model.py @@ -18,7 +18,7 @@ import pyvista as pv from physiomotion4d.contour_tools import ContourTools -from physiomotion4d.workflow_convert_ct_to_vtk import WorkflowConvertCTToVTK +from physiomotion4d.workflow_convert_image_to_vtk import WorkflowConvertImageToVTK from physiomotion4d.workflow_create_statistical_model import ( WorkflowCreateStatisticalModel, ) @@ -116,8 +116,8 @@ def create_meshes( print(f"Segmenting {case_prefix} from {phase_files[0].name}") image = itk.imread(str(phase_files[0])) - workflow = WorkflowConvertCTToVTK( - segmentation_method="total_segmentator", + workflow = WorkflowConvertImageToVTK( + segmentation_method="ChestTotalSegmentator", log_level=log_level, ) result = workflow.run_workflow( From 5af4a848c63f8440847c0d9ab4c829392d2a7fb4 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Sun, 17 May 2026 21:04:15 -0400 Subject: [PATCH 4/4] BUG: improvements from AI reviews --- CLAUDE.md | 3 +- docs/API_MAP.md | 24 ++++----- docs/cli_scripts/byod_tutorials.rst | 10 ++-- src/physiomotion4d/register_images_icon.py | 54 ++++++++++++++++++- .../workflow_convert_image_to_usd.py | 43 ++++++++++++--- ...rkflow_fit_statistical_model_to_patient.py | 5 ++ 6 files changed, 111 insertions(+), 28 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index aca3523..b7c040c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -16,7 +16,7 @@ We are developing open-source code for scientific AI libraries. Leverage GPU-acc 5) testing ## Behavior Guidelines -1) Don't assume. Don't hide confuction. Surface tradeoffs. +1) Don't assume. Don't hide confusion. Surface tradeoffs. 2) Minimum code that solves the problem. Nothing speculative. 3) Touch only what you must. Clean up only your own mess. 4) Define success criteria. Loop until verified. @@ -103,6 +103,7 @@ Before editing any code: 8. Implement in small, reviewable diffs. 9. Update docstrings and tests for every changed public method. 10. Call out breaking changes explicitly. +11. Do not commit changes or make pull requests unless specifically told to do so. Breaking changes are acceptable. Backward-compatibility shims are not. diff --git a/docs/API_MAP.md b/docs/API_MAP.md index b3493a4..0a3e8b0 100644 --- a/docs/API_MAP.md +++ b/docs/API_MAP.md @@ -210,7 +210,7 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - `def set_mass_preservation(self, enable)` (line 125): Enable or disable mass preservation constraint. - `def preprocess(self, image, modality='ct')` (line 142): Preprocess the image for ICON registration. - `def registration_method(self, moving_image, moving_mask=None, moving_labelmap=None, moving_image_pre=None, initial_forward_transform=None)` (line 162): Register moving image to fixed image using ICON registration algorithm. - - `def finetune(self, image_pairs, output_model_filename, mask_pairs=None, epochs=1, learning_rate=DEFAULT_FINETUNE_LEARNING_RATE)` (line 336): Fine-tune the ICON network on a cohort of image pairs. + - `def finetune(self, image_pairs, output_model_filename, mask_pairs=None, epochs=1, learning_rate=DEFAULT_FINETUNE_LEARNING_RATE)` (line 386): Fine-tune the ICON network on a cohort of image pairs. ## src/physiomotion4d/register_models_distance_maps.py @@ -460,17 +460,17 @@ _Re-run `py utils/generate_api_map.py` whenever public APIs change._ - **class WorkflowFitStatisticalModelToPatient** (line 56): Register anatomical models using multi-stage ICP, mask-based, and image-based - `def __init__(self, template_model, patient_models=None, patient_image=None, segmentation_method='HeartSimpleware', log_level=logging.INFO)` (line 135): Initialize the model-to-image-and-model registration pipeline. - - `def set_mask_dilation_mm(self, mask_dilation_mm)` (line 356): Set mask dilation amount for auto-generated masks. - - `def set_roi_dilation_mm(self, roi_dilation_mm)` (line 365): Set ROI mask dilation amount. - - `def set_use_pca_registration(self, use_pca_registration, pca_model=None, pca_number_of_modes=0, pca_uses_surface=True)` (line 374): Set whether to use PCA-based registration and provide the PCA model. - - `def set_use_mask_to_mask_registration(self, use_mask_to_mask_registration)` (line 409): Set whether to use mask-to-mask registration. - - `def set_use_mask_to_image_registration(self, use_mask_to_image_registration, template_labelmap=None, template_labelmap_organ_mesh_ids=None, template_labelmap_organ_extra_ids=None, template_labelmap_background_ids=None)` (line 420): Set whether to use mask-to-image registration. - - `def register_model_to_model_icp(self)` (line 494): Perform ICP alignment of template model to patient model. - - `def register_model_to_model_pca(self)` (line 552): Perform PCA-based registration after ICP alignment. - - `def register_mask_to_mask(self, use_icon_refinement=False)` (line 678): Perform mask-based deformable registration of model to patient model. - - `def register_labelmap_to_image(self, use_icon_refinement=False)` (line 746): Perform labelmap-to-image refinement. - - `def transform_model(self, base_model=None)` (line 866): Apply registration transforms to the model. - - `def run_workflow(self, use_icon_registration_refinement=False)` (line 931): Execute the complete multi-stage registration workflow. + - `def set_mask_dilation_mm(self, mask_dilation_mm)` (line 361): Set mask dilation amount for auto-generated masks. + - `def set_roi_dilation_mm(self, roi_dilation_mm)` (line 370): Set ROI mask dilation amount. + - `def set_use_pca_registration(self, use_pca_registration, pca_model=None, pca_number_of_modes=0, pca_uses_surface=True)` (line 379): Set whether to use PCA-based registration and provide the PCA model. + - `def set_use_mask_to_mask_registration(self, use_mask_to_mask_registration)` (line 414): Set whether to use mask-to-mask registration. + - `def set_use_mask_to_image_registration(self, use_mask_to_image_registration, template_labelmap=None, template_labelmap_organ_mesh_ids=None, template_labelmap_organ_extra_ids=None, template_labelmap_background_ids=None)` (line 425): Set whether to use mask-to-image registration. + - `def register_model_to_model_icp(self)` (line 499): Perform ICP alignment of template model to patient model. + - `def register_model_to_model_pca(self)` (line 557): Perform PCA-based registration after ICP alignment. + - `def register_mask_to_mask(self, use_icon_refinement=False)` (line 683): Perform mask-based deformable registration of model to patient model. + - `def register_labelmap_to_image(self, use_icon_refinement=False)` (line 751): Perform labelmap-to-image refinement. + - `def transform_model(self, base_model=None)` (line 871): Apply registration transforms to the model. + - `def run_workflow(self, use_icon_registration_refinement=False)` (line 936): Execute the complete multi-stage registration workflow. ## src/physiomotion4d/workflow_reconstruct_highres_4d_ct.py diff --git a/docs/cli_scripts/byod_tutorials.rst b/docs/cli_scripts/byod_tutorials.rst index fa6a122..22881f2 100644 --- a/docs/cli_scripts/byod_tutorials.rst +++ b/docs/cli_scripts/byod_tutorials.rst @@ -34,7 +34,7 @@ Verify that both relevant CLI entry-points are available after installation: .. code-block:: bash - physiomotion4d-convert-heart-gated-ct-to-usd --help + physiomotion4d-convert-image-to-usd --help physiomotion4d-convert-vtk-to-usd --help See :doc:`/installation` for prerequisites, CUDA version requirements, and @@ -56,7 +56,7 @@ Pass a single ``.nii.gz`` file to produce a static USD scene. .. code-block:: bash - physiomotion4d-convert-heart-gated-ct-to-usd \ + physiomotion4d-convert-image-to-usd \ patient_ct.nii.gz \ --output patient_heart.usd @@ -66,7 +66,7 @@ Pass a single ``.nii.gz`` file to produce a static USD scene. import physiomotion4d as pm4d - wf = pm4d.WorkflowConvertHeartGatedCTToUSD() + wf = pm4d.WorkflowConvertImageToUSD() wf.input_image = "patient_ct.nii.gz" wf.output_file = "patient_heart.usd" wf.run_workflow() @@ -88,7 +88,7 @@ animated USD scene. Use ``--fps`` to control playback rate and .. code-block:: bash - physiomotion4d-convert-heart-gated-ct-to-usd \ + physiomotion4d-convert-image-to-usd \ phase_*.nii.gz \ --output heart_animated.usd \ --fps 25 \ @@ -101,7 +101,7 @@ animated USD scene. Use ``--fps`` to control playback rate and import glob import physiomotion4d as pm4d - wf = pm4d.WorkflowConvertHeartGatedCTToUSD() + wf = pm4d.WorkflowConvertImageToUSD() wf.input_images = sorted(glob.glob("phase_*.nii.gz")) wf.output_file = "heart_animated.usd" wf.fps = 25 diff --git a/src/physiomotion4d/register_images_icon.py b/src/physiomotion4d/register_images_icon.py index 31072cf..8f60a2a 100644 --- a/src/physiomotion4d/register_images_icon.py +++ b/src/physiomotion4d/register_images_icon.py @@ -313,7 +313,34 @@ def _image_to_resized_tensor( """Convert an itk image to a torch tensor resized to the net's input grid. Mirrors the trilinear preprocessing path used by - ``icon_registration.itk_wrapper.register_pair``. + ``icon_registration.itk_wrapper.register_pair`` exactly. + + Axis ordering: + - Input ``image`` is a scalar (single-channel) 3D ``itk.Image`` with + ITK world-axis order (X, Y, Z). ``np.array(image)`` returns a + C-contiguous array with axes **reversed** to (Z, Y, X) — ITK's + standard numpy view. + - ``torch.Tensor(arr)`` casts to ``float32`` (PyTorch's + ``FloatTensor`` constructor) regardless of the source dtype. + - Indexing with ``[None, None]`` prepends batch and channel + singleton axes, producing shape ``(1, 1, Z, Y, X)``. This is + PyTorch's NCDHW layout where ``D=Z``, ``H=Y``, ``W=X``. + - ``shape`` is ``self.net.identity_map.shape`` (5D, NCDHW); the + target spatial size is ``shape[2:] = (D_out, H_out, W_out)``. + - Return shape: ``(1, 1, D_out, H_out, W_out)``, float32, + C-contiguous on ``icon.config.device``. + + Notes: + - Single-channel scalar inputs only. Vector/multi-channel + ``itk.Image`` would yield ``(Z, Y, X, C)`` from ``np.array`` and + break the assumed NCDHW layout — not supported here, matching + ICON's own preprocessing. + - No explicit time axis: 4D series must be split into 3D + timepoints by the caller; pairs are processed one volume at a + time. + - No transpose is performed; the (Z, Y, X) numpy ordering is + consumed directly as (D, H, W). Voxel values, not world + coordinates, drive the trilinear resample. """ arr = np.array(image) tensor = torch.Tensor(arr).to(icon.config.device)[None, None] @@ -327,7 +354,30 @@ def _mask_to_resized_tensor( """Convert an itk mask image to a torch tensor resized via nearest-neighbor. Mirrors the mask preprocessing used by - ``icon_registration.itk_wrapper.register_pair_with_mask``. + ``icon_registration.itk_wrapper.register_pair_with_mask`` exactly. + + Axis ordering: + - Input ``mask`` is a scalar (single-channel) 3D ``itk.Image`` + (typically ``uint8``/short labels) with ITK world-axis order + (X, Y, Z). ``np.array(mask)`` returns a C-contiguous array with + axes **reversed** to (Z, Y, X). + - ``torch.Tensor(arr)`` casts to ``float32`` (label values become + integral-valued floats; nearest-neighbor resampling preserves + them). + - ``[None, None]`` prepends batch and channel singletons → + ``(1, 1, Z, Y, X)`` in NCDHW (``D=Z``, ``H=Y``, ``W=X``). + - Target spatial size is ``shape[2:] = (D_out, H_out, W_out)`` + from ``self.net.identity_map.shape``. + - Return shape: ``(1, 1, D_out, H_out, W_out)``, float32, + C-contiguous on ``icon.config.device``. + + Notes: + - Single-channel mask inputs only; multi-label masks are encoded + as scalar integer values, not channels. Vector ``itk.Image`` + inputs are not supported. + - No time axis: per-volume 3D processing. + - Resampling uses ``mode='nearest'`` (no ``align_corners``) so + label identities are preserved. """ arr = np.array(mask) tensor = torch.Tensor(arr).to(icon.config.device)[None, None] diff --git a/src/physiomotion4d/workflow_convert_image_to_usd.py b/src/physiomotion4d/workflow_convert_image_to_usd.py index e19dfce..aa0e77b 100644 --- a/src/physiomotion4d/workflow_convert_image_to_usd.py +++ b/src/physiomotion4d/workflow_convert_image_to_usd.py @@ -281,6 +281,26 @@ def _load_time_series(self) -> None: self.log_info("Loaded %d time points", self._num_time_points) + def _optional_mask(self, key: str) -> itk.Image: + """Return ``self._fixed_segmentation[key]`` or a zero mask if absent. + + Segmenters expose only the anatomy groups they actually produce + (see ``SegmentAnatomyBase.create_anatomy_group_masks`` — + ``SegmentHeartSimpleware`` omits ``lung``/``bone``/``soft_tissue``/ + ``contrast``, while ``SegmentChestTotalSegmentator`` provides them). + Treating missing groups as empty (uint8 zeros matching + ``self._fixed_image``) lets downstream mask arithmetic run + uniformly for any ``self.segmenter`` choice. + """ + assert self._fixed_segmentation is not None, "Fixed segmentation must be set" + assert self._fixed_image is not None, "Fixed image must be set" + if key in self._fixed_segmentation: + return self._fixed_segmentation[key] + zeros = np.zeros(itk.array_from_image(self._fixed_image).shape, dtype=np.uint8) + empty = itk.GetImageFromArray(zeros) + empty.CopyInformation(self._fixed_image) + return empty + def _segment_and_register_frames(self) -> None: """Segment each frame and register to reference image.""" self.log_info("Segmenting and registering frames...") @@ -292,15 +312,18 @@ def _segment_and_register_frames(self) -> None: self._fixed_image, contrast_enhanced_study=self.contrast_enhanced ) - # Create combined masks for registration + # Create combined masks for registration. Optional groups + # (lung/bone/contrast) are absent from segmenters that do not produce + # them (e.g., SegmentHeartSimpleware) — fall back to empty masks so + # the static/dynamic-mask sums below remain well-defined. assert self._fixed_segmentation is not None, "Fixed segmentation must be set" labelmap_mask = self._fixed_segmentation["labelmap"] - lung_mask = self._fixed_segmentation["lung"] heart_mask = self._fixed_segmentation["heart"] major_vessels_mask = self._fixed_segmentation["major_vessels"] - bone_mask = self._fixed_segmentation["bone"] other_mask = self._fixed_segmentation["other"] - contrast_mask = self._fixed_segmentation["contrast"] + lung_mask = self._optional_mask("lung") + bone_mask = self._optional_mask("bone") + contrast_mask = self._optional_mask("contrast") self._write_image_if_enabled( labelmap_mask, "fixed_image_mask.mha", @@ -433,15 +456,19 @@ def _generate_reference_contours(self) -> None: """Generate contour meshes from reference segmentation.""" self.log_info("Generating reference contours...") + # Optional groups (lung/bone/soft_tissue/contrast) are absent from + # segmenters that do not produce them (e.g., SegmentHeartSimpleware) + # — fall back to empty masks so the static/dynamic-anatomy sums below + # remain well-defined. assert self._fixed_segmentation is not None, "Fixed segmentation must be set" labelmap_image = self._fixed_segmentation["labelmap"] - lung_mask = self._fixed_segmentation["lung"] heart_mask = self._fixed_segmentation["heart"] major_vessels_mask = self._fixed_segmentation["major_vessels"] - bone_mask = self._fixed_segmentation["bone"] - soft_tissue_mask = self._fixed_segmentation["soft_tissue"] other_mask = self._fixed_segmentation["other"] - contrast_mask = self._fixed_segmentation["contrast"] + lung_mask = self._optional_mask("lung") + bone_mask = self._optional_mask("bone") + soft_tissue_mask = self._optional_mask("soft_tissue") + contrast_mask = self._optional_mask("contrast") # Generate all anatomy contours all_contours = self.contour_tools.extract_contours(labelmap_image) diff --git a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py index 5cacd62..72945bc 100644 --- a/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiomotion4d/workflow_fit_statistical_model_to_patient.py @@ -149,6 +149,11 @@ def __init__( patient_image: Optional patient image providing the target coordinate frame. If None, a reference image is created from the patient model surface via create_reference_image (contour_tools). + segmentation_method: Segmentation backend used by + WorkflowConvertImageToVTK when patient_models is None and + patient_image is provided. One of ``'HeartSimpleware'`` (default) + or ``'ChestTotalSegmentator'``. Ignored when patient_models is + supplied. log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO """