ENH: Demo for combining heart and lung motion#90
Conversation
Add Tutorial 11 that fuses cardiac motion (trained PhysicsNeMo MeshGraphNet applied to the patient heart) with respiratory motion (Tutorial 1 lung transforms) into one animated 4D USD, split by anatomy label and painted with per-organ OmniSurface materials. Add the Heart_and_Lungs_Motion experiment scripts behind it. usd_anatomy_tools: expand DEFAULT_RENDER_PARAMS with organ-level materials (skin, airway, vein/artery, muscle, fat, cartilage, endocrine organs, oxygenation-coded heart chambers) and switch enhance_meshes override matching to longest-substring-wins so e.g. "kidney" covers kidney_left/right. segment_chest_total_segmentator: move sacrum (label 25) into the correct group; stop overwriting body labels so body_skin (133) survives. Tutorials 05/09/10: evaluate the resumed run's output directory instead of OUTPUT_DIR, add missing casts/asserts/type hints for mypy. CI: run Ruff through pre-commit so its version is pinned in one place; bump ruff-pre-commit to v0.15.20. Add .claude-plugin marketplace/plugin manifests and .cursor rules; restructure CLAUDE.md behavior guidelines into numbered sections.
WalkthroughThe PR adds repository AI guidance and plugin metadata, updates CI and workflow return contracts, refines anatomy processing, extends PhysicsNeMo inference outputs, and introduces CT, cardiac, and combined cardiac-respiratory motion tutorials and experiments. ChangesRepository workflows and tooling
Inference and motion pipelines
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant WorkflowInferPhysicsNeMo
participant TransformTools
participant WorkflowConvertVTKToUSD
participant USDAnatomyTools
WorkflowInferPhysicsNeMo->>TransformTools: produce cardiac deformation fields
TransformTools->>WorkflowInferPhysicsNeMo: apply cardiac and respiratory transforms
WorkflowInferPhysicsNeMo->>WorkflowConvertVTKToUSD: provide ordered surface frames
WorkflowConvertVTKToUSD->>USDAnatomyTools: enhance and split USD meshes
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
src/physiotwin4d/usd_anatomy_tools.py (1)
834-839: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
GROUP_RENDER_KEYSfor consistent override key resolution.Relying on
taxonomy.group_names()could cause discrepancies if the taxonomy is missing a group (such ascontrast). If a group exists inself.render_paramsbut not in the taxonomy, it would incorrectly be treated as an organ-level override key. Using the canonicalGROUP_RENDER_KEYSensures consistency with_resolve_render_paramsand prevents subtle matching bugs.♻️ Proposed refactor
- group_names = set(taxonomy.group_names()) override_keys = sorted( - (k for k in self.render_params if k not in group_names and k != "other"), + (k for k in self.render_params if k not in GROUP_RENDER_KEYS), key=len, reverse=True, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/usd_anatomy_tools.py` around lines 834 - 839, Update the override-key filtering near `group_names` to use the canonical `GROUP_RENDER_KEYS` set instead of `taxonomy.group_names()`. Preserve the existing exclusions for `"other"` and the length-descending sort so resolution remains consistent with `_resolve_render_params`.src/physiotwin4d/segment_chest_total_segmentator.py (1)
383-384: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove commented-out code.
As per coding guidelines, remove unused code made obsolete by your changes.
♻️ Proposed refactor
- # mask = final_arr > 0 - # labelmap_arr_body[mask] = 0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/physiotwin4d/segment_chest_total_segmentator.py` around lines 383 - 384, Remove the obsolete commented-out mask assignment lines near the final label-map processing in the chest segmentation flow. Do not alter the surrounding active logic.Source: Coding guidelines
experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py (1)
71-113: 📐 Maintainability & Code Quality | 🔵 Trivial
_ensure_mgn_inference_assetsduplicated verbatim intutorials/tutorial_11_heart_and_lung_motion.py(L95-136).Same asset-finalization logic (checkpoint copy, mean-surface extraction, shared graph tensors) is copy-pasted across two files. See the consolidated comment below for the proposed extraction.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py` around lines 71 - 113, Extract the shared asset-finalization logic from _ensure_mgn_inference_assets into a reusable helper, then update both _ensure_mgn_inference_assets implementations in the heart-and-lungs scripts to call it instead of duplicating checkpoint, surface, and graph-tensor generation. Preserve the existing idempotent behavior and inputs, including the epoch checkpoint and PCA mean volume handling.tutorials/tutorial_11_heart_and_lung_motion.py (1)
95-185: 📐 Maintainability & Code Quality | 🔵 TrivialThree helper functions duplicate the experiment scripts' identical logic.
_ensure_mgn_inference_assets,_smoothed_cardiac_transform, and_condition_surfaceeach restate logic already present inexperiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.pyandexperiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py. See the consolidated comment for the proposed shared-module extraction.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tutorials/tutorial_11_heart_and_lung_motion.py` around lines 95 - 185, The helper functions _ensure_mgn_inference_assets, _smoothed_cardiac_transform, and _condition_surface duplicate logic from the two Heart_and_Lungs_Motion experiment scripts. Extract their shared implementations into the proposed common module, update both experiment scripts and this tutorial to import and reuse those shared symbols, and remove the duplicate local definitions while preserving their current behavior.experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py (1)
92-128: 📐 Maintainability & Code Quality | 🔵 Trivial
_condition_surfaceand_smoothed_cardiac_transformduplicate the logic intutorials/tutorial_11_heart_and_lung_motion.py(L159-185, L139-156).Same decimate/smooth and field-wrap/smooth steps, with only minor signature differences (module-level
_transform_toolshere vs. an explicit parameter in the tutorial). See consolidated comment.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py` around lines 92 - 128, Remove the duplicated implementations of _condition_surface and _smoothed_cardiac_transform, and reuse the corresponding shared helpers from tutorials/tutorial_11_heart_and_lung_motion.py or the consolidated utility they expose. Preserve the existing decimation, smoothing, field conversion, and transform-smoothing behavior while adapting only the differing transform-tools dependency.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tutorials/tutorial_02_lung_ct_to_vtk.py`:
- Around line 97-111: Filter out None-valued entries from result["surfaces"]
before passing them to ContourTools.save_combined_surface and the conditional
ContourTools.save_surfaces call. Apply the same filtering to any surface
collection used for saving, while preserving result["label_surfaces"] behavior
unless it also requires filtering based on its contents.
In `@tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py`:
- Around line 102-113: Update the three itk.imwrite calls in the tutorial
flow—saving patient_image, patient_labelmap, and heart_labelmap—to pass
compression=True, preserving their existing output paths.
---
Nitpick comments:
In `@experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py`:
- Around line 71-113: Extract the shared asset-finalization logic from
_ensure_mgn_inference_assets into a reusable helper, then update both
_ensure_mgn_inference_assets implementations in the heart-and-lungs scripts to
call it instead of duplicating checkpoint, surface, and graph-tensor generation.
Preserve the existing idempotent behavior and inputs, including the epoch
checkpoint and PCA mean volume handling.
In `@experiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.py`:
- Around line 92-128: Remove the duplicated implementations of
_condition_surface and _smoothed_cardiac_transform, and reuse the corresponding
shared helpers from tutorials/tutorial_11_heart_and_lung_motion.py or the
consolidated utility they expose. Preserve the existing decimation, smoothing,
field conversion, and transform-smoothing behavior while adapting only the
differing transform-tools dependency.
In `@src/physiotwin4d/segment_chest_total_segmentator.py`:
- Around line 383-384: Remove the obsolete commented-out mask assignment lines
near the final label-map processing in the chest segmentation flow. Do not alter
the surrounding active logic.
In `@src/physiotwin4d/usd_anatomy_tools.py`:
- Around line 834-839: Update the override-key filtering near `group_names` to
use the canonical `GROUP_RENDER_KEYS` set instead of `taxonomy.group_names()`.
Preserve the existing exclusions for `"other"` and the length-descending sort so
resolution remains consistent with `_resolve_render_params`.
In `@tutorials/tutorial_11_heart_and_lung_motion.py`:
- Around line 95-185: The helper functions _ensure_mgn_inference_assets,
_smoothed_cardiac_transform, and _condition_surface duplicate logic from the two
Heart_and_Lungs_Motion experiment scripts. Extract their shared implementations
into the proposed common module, update both experiment scripts and this
tutorial to import and reuse those shared symbols, and remove the duplicate
local definitions while preserving their current behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8238208c-f571-4fa5-b241-6ffa0cfd128c
📒 Files selected for processing (25)
.claude-plugin/marketplace.json.claude-plugin/plugin.json.cursor/rules/karpathy-guidelines.mdc.cursor/rules/project-standards.mdc2.github/workflows/ci.yml.gitignore.pre-commit-config.yamlCLAUDE.mddocs/api/workflows.rstdocs/tutorials.rstexperiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.pyexperiments/Heart_and_Lungs_Motion/1-heart_and_lungs_combined.pysrc/physiotwin4d/cli/convert_vtk_to_usd.pysrc/physiotwin4d/segment_chest_total_segmentator.pysrc/physiotwin4d/usd_anatomy_tools.pysrc/physiotwin4d/workflow_convert_vtk_to_usd.pysrc/physiotwin4d/workflow_infer_physicsnemo.pytutorials/tutorial_02_lung_ct_to_vtk.pytutorials/tutorial_03_heart_vtk_to_usd.pytutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.pytutorials/tutorial_09_byod_train_physicsnemo_mgn.pytutorials/tutorial_09_byod_train_physicsnemo_mlp.pytutorials/tutorial_10_byod_eval_physicsnemo_mgn.pytutorials/tutorial_10_byod_eval_physicsnemo_mlp.pytutorials/tutorial_11_heart_and_lung_motion.py
💤 Files with no reviewable changes (1)
- .gitignore
| surface_file = Path( | ||
| ContourTools.save_combined_surface( | ||
| result["surfaces"], | ||
| str(output_dir), | ||
| prefix="patient", | ||
| ) | ||
| ) | ||
| if save_group_surfaces: | ||
| ContourTools.save_surfaces( | ||
| result["surfaces"], str(output_dir), prefix="patient" | ||
| ) | ||
| if save_label_surfaces: | ||
| ContourTools.save_surfaces( | ||
| result["label_surfaces"], str(output_dir), prefix="patient" | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Filter None values before saving surfaces.
The list comprehension at line 137 (if surface is not None) implies that result["surfaces"] can contain None values (e.g., if a particular anatomical group was not found). If None values are present, passing the unfiltered dictionaries to ContourTools.save_combined_surface() and ContourTools.save_surfaces() will cause a crash (e.g., AttributeError: 'NoneType' object has no attribute 'save').
Filter out None values before passing the dictionaries to the save functions to ensure stability.
🐛 Proposed fix to filter invalid surfaces
- surface_file = Path(
- ContourTools.save_combined_surface(
- result["surfaces"],
- str(output_dir),
- prefix="patient",
- )
- )
- if save_group_surfaces:
- ContourTools.save_surfaces(
- result["surfaces"], str(output_dir), prefix="patient"
- )
- if save_label_surfaces:
- ContourTools.save_surfaces(
- result["label_surfaces"], str(output_dir), prefix="patient"
- )
+ valid_surfaces = {k: v for k, v in result["surfaces"].items() if v is not None}
+ valid_label_surfaces = {
+ k: v for k, v in result.get("label_surfaces", {}).items() if v is not None
+ }
+
+ surface_file = Path(
+ ContourTools.save_combined_surface(
+ valid_surfaces,
+ str(output_dir),
+ prefix="patient",
+ )
+ )
+ if save_group_surfaces:
+ ContourTools.save_surfaces(
+ valid_surfaces, str(output_dir), prefix="patient"
+ )
+ if save_label_surfaces:
+ ContourTools.save_surfaces(
+ valid_label_surfaces, str(output_dir), prefix="patient"
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| surface_file = Path( | |
| ContourTools.save_combined_surface( | |
| result["surfaces"], | |
| str(output_dir), | |
| prefix="patient", | |
| ) | |
| ) | |
| if save_group_surfaces: | |
| ContourTools.save_surfaces( | |
| result["surfaces"], str(output_dir), prefix="patient" | |
| ) | |
| if save_label_surfaces: | |
| ContourTools.save_surfaces( | |
| result["label_surfaces"], str(output_dir), prefix="patient" | |
| ) | |
| valid_surfaces = {k: v for k, v in result["surfaces"].items() if v is not None} | |
| valid_label_surfaces = { | |
| k: v for k, v in result.get("label_surfaces", {}).items() if v is not None | |
| } | |
| surface_file = Path( | |
| ContourTools.save_combined_surface( | |
| valid_surfaces, | |
| str(output_dir), | |
| prefix="patient", | |
| ) | |
| ) | |
| if save_group_surfaces: | |
| ContourTools.save_surfaces( | |
| valid_surfaces, str(output_dir), prefix="patient" | |
| ) | |
| if save_label_surfaces: | |
| ContourTools.save_surfaces( | |
| valid_label_surfaces, str(output_dir), prefix="patient" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tutorials/tutorial_02_lung_ct_to_vtk.py` around lines 97 - 111, Filter out
None-valued entries from result["surfaces"] before passing them to
ContourTools.save_combined_surface and the conditional
ContourTools.save_surfaces call. Apply the same filtering to any surface
collection used for saving, while preserving result["label_surfaces"] behavior
unless it also requires filtering based on its contents.
| itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz") | ||
|
|
||
| heart_labelmap = segmentation_result["heart"] | ||
| itk.imwrite(heart_labelmap, output_dir / "heart_labelmap.nii.gz") | ||
| segmentation_result = segmentation_method.segment(patient_image) | ||
| patient_labelmap = segmentation_result["labelmap"] | ||
| itk.imwrite( | ||
| patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz" | ||
| ) | ||
|
|
||
| heart_labelmap = segmentation_result["heart"] | ||
| itk.imwrite( | ||
| heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz" | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Use compression=True when saving ITK images.
The itk.imwrite calls are missing the compression argument. As per coding guidelines, store ITK images with itk.imwrite(..., compression=True).
🐛 Proposed fix
- itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz")
+ itk.imwrite(
+ patient_image,
+ output_dir / f"{project_name}_patient_image.nii.gz",
+ compression=True,
+ )
segmentation_result = segmentation_method.segment(patient_image)
patient_labelmap = segmentation_result["labelmap"]
- itk.imwrite(
- patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz"
- )
+ itk.imwrite(
+ patient_labelmap,
+ output_dir / f"{project_name}_patient_labelmap.nii.gz",
+ compression=True,
+ )
heart_labelmap = segmentation_result["heart"]
- itk.imwrite(
- heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz"
- )
+ itk.imwrite(
+ heart_labelmap,
+ output_dir / f"{project_name}_heart_labelmap.nii.gz",
+ compression=True,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| itk.imwrite(patient_image, output_dir / f"{project_name}_patient_image.nii.gz") | |
| heart_labelmap = segmentation_result["heart"] | |
| itk.imwrite(heart_labelmap, output_dir / "heart_labelmap.nii.gz") | |
| segmentation_result = segmentation_method.segment(patient_image) | |
| patient_labelmap = segmentation_result["labelmap"] | |
| itk.imwrite( | |
| patient_labelmap, output_dir / f"{project_name}_patient_labelmap.nii.gz" | |
| ) | |
| heart_labelmap = segmentation_result["heart"] | |
| itk.imwrite( | |
| heart_labelmap, output_dir / f"{project_name}_heart_labelmap.nii.gz" | |
| ) | |
| itk.imwrite( | |
| patient_image, | |
| output_dir / f"{project_name}_patient_image.nii.gz", | |
| compression=True, | |
| ) | |
| segmentation_result = segmentation_method.segment(patient_image) | |
| patient_labelmap = segmentation_result["labelmap"] | |
| itk.imwrite( | |
| patient_labelmap, | |
| output_dir / f"{project_name}_patient_labelmap.nii.gz", | |
| compression=True, | |
| ) | |
| heart_labelmap = segmentation_result["heart"] | |
| itk.imwrite( | |
| heart_labelmap, | |
| output_dir / f"{project_name}_heart_labelmap.nii.gz", | |
| compression=True, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tutorials/tutorial_05_heart_to_lung_fit_statistical_model_to_patient.py`
around lines 102 - 113, Update the three itk.imwrite calls in the tutorial
flow—saving patient_image, patient_labelmap, and heart_labelmap—to pass
compression=True, preserving their existing output paths.
Source: Coding guidelines
Summary by CodeRabbit