ENH: Improve fit for PCA and Labelmap_to_Labelmap registration#69
Conversation
WalkthroughThe PR extends the statistical model-to-patient workflow with optional labelmap inputs and exclusion IDs, adds isotropic image resampling, changes PCA and refinement registration inputs, updates distance-map handling, and adjusts supporting tests, experiment wiring, utility fallback behavior, and CI. ChangesRegistration workflow and supporting updates
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 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.
Pull request overview
This PR aims to improve the statistical-model fitting workflow by refining PCA-based registration inputs and strengthening labelmap-based registration stages (labelmap-to-labelmap and labelmap-to-image), with some supporting utility refactors.
Changes:
- Refactors isotropic-resampling into
ImageTools.make_isotropic_image()and uses it in the fit workflow when patient spacing is anisotropic. - Improves labelmap propagation/selection for labelmap-to-image refinement and adds support for supplying external template/patient labelmaps to the workflow.
- Adjusts distance-map registration behavior (metrics, masking behavior) and Greedy’s metric-grid downsampling rationale/usage.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_workflow_fit_statistical_model_to_patient.py | Updates tests to use renamed workflow flags (use_l2l_registration, use_l2i_registration). |
| src/physiomotion4d/workflow_fit_statistical_model_to_patient.py | Adds optional labelmap inputs, uses isotropic resampling via ImageTools, and changes PCA/L2I staging/propagation behavior. |
| src/physiomotion4d/transform_tools.py | Broadens CuPy import failure handling to include OSError. |
| src/physiomotion4d/register_models_pca.py | Changes distance-map generation options and introduces new invalid-point handling in the metric. |
| src/physiomotion4d/register_models_distance_maps.py | Alters distance-map preprocessing and registration configuration (Greedy metric, ICON masking behavior). |
| src/physiomotion4d/register_images_greedy.py | Updates comments and applies metric downsampling scale consistently via _metric_downsample_scale(). |
| src/physiomotion4d/image_tools.py | Introduces make_isotropic_image() utility method. |
| src/physiomotion4d/contour_tools.py | Changes labelmap array dtype from uint8 to uint16. |
| experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py | Enables labelmap-to-labelmap registration in the experiment script. |
Comments suppressed due to low confidence (2)
src/physiomotion4d/register_models_pca.py:161
- Unconditional itk.imwrite() in the constructor will write files to the current working directory on every run, which is a surprising side effect for library code and can break read-only environments. Gate this behind a debug log level (or remove entirely).
self.pca_number_of_modes: int = pca_number_of_modes
src/physiomotion4d/register_models_distance_maps.py:354
- Commented-out mask handling makes the ICON stage behavior unclear and leaves dead code in the method. Either keep mask support enabled or remove the commented lines and update behavior intentionally (optionally via a parameter/flag).
# self.reference_image,
# interpolation_method="nearest",
# )
# Configure and run ICON
self.registrar_ICON.set_number_of_iterations(icon_iterations)
self.registrar_ICON.set_fixed_image(self.fixed_distance_map_image)
# self.registrar_ICON.set_fixed_mask(self.fixed_mask_image)
result_ICON = self.registrar_ICON.register(
moving_image=moving_distance_map_affine_transformed,
# moving_mask=moving_mask_affine_transformed,
)
forward_transform_ICON = result_ICON["forward_transform"]
inverse_transform_ICON = result_ICON["inverse_transform"]
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #69 +/- ##
==========================================
- Coverage 32.45% 32.38% -0.08%
==========================================
Files 53 53
Lines 7244 7257 +13
==========================================
- Hits 2351 2350 -1
- Misses 4893 4907 +14
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| # Python 3.11 is required because the [physicsnemo] extra pulls in | ||
| # nvidia-physicsnemo, which requires Python >= 3.11. | ||
| run: | | ||
| & "C:\Program Files\Python311\python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" | ||
| & "C:\Users\saylward\AppData\Local\Programs\Python\Python311/python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" | ||
| echo "$env:RUNNER_TEMP\physiomotion4d-venv\Scripts" >> $env:GITHUB_PATH |
| patient_models: List of patient-specific models extracted from imaging | ||
| data. Typically 3 models for cardiac applications: LV, myocardium, RV. | ||
| 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). |
| patient_models: list[pv.DataSet] | None = None, | ||
| patient_image: Optional[itk.Image] = None, | ||
| patient_labelmap: Optional[itk.Image] = None, | ||
| template_labelmap: Optional[itk.Image] = None, | ||
| labelmap_interior_object_ids: Optional[list] = None, | ||
| segmentation_method: str = "HeartSimplewareTrimmedBranches", | ||
| log_level: int | str = logging.INFO, |
| self.labelmap_interior_object_ids: Optional[list] = labelmap_interior_object_ids | ||
|
|
| def set_mask_dilation_mm(self, mask_dilation_mm: float) -> None: | ||
| """Set dilation amount for binary registration masks. | ||
|
|
||
| Args: | ||
| mask_dilation_mm: Dilation amount in millimeters for binary registration | ||
| mask generation. Default: 25mm | ||
| mask generation. Default: 10mm | ||
| """ | ||
| self.mask_dilation_mm = mask_dilation_mm |
| if n_invalid_points >= 0.05 * n_valid_points: | ||
| self.log_warning( | ||
| f"{n_invalid_points} of {n_valid_points + n_invalid_points} mapped outside of image. Rejecting." | ||
| ) | ||
| return self._fixed_distance_map_max_distance |
| # moving_mask_affine_transformed = self.transform_tools.transform_image( | ||
| # self.moving_mask_image, | ||
| # forward_transform_Greedy, | ||
| # self.reference_image, | ||
| # interpolation_method="nearest", | ||
| # ) |
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/physiomotion4d/contour_tools.py (1)
261-272: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winGuard the new
uint16label capacity.
i + 1can exceednp.uint16capacity, which would corrupt or reject anatomy labels. Add a clear bounds check before filling the array.Proposed fix
+ if len(meshes) > np.iinfo(np.uint16).max: + raise ValueError( + "create_labelmap_from_meshes supports at most " + f"{np.iinfo(np.uint16).max} labels with uint16 output" + ) labelmap_arr = np.zeros(As per coding guidelines, masks/labelmaps must use integer labels with consistent anatomy group IDs across segmenters.
🤖 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/physiomotion4d/contour_tools.py` around lines 261 - 272, The labelmap in contour_tools.py is now using np.uint16, but the assignment in the mesh loop can still overflow when i + 1 exceeds the available label range. Add a bounds check in the labelmap-building logic around create_mask_from_mesh/labelmap_arr assignment so the code fails fast or skips with a clear error before writing an out-of-range label, keeping anatomy group IDs consistent.Source: Coding guidelines
🧹 Nitpick comments (2)
src/physiomotion4d/workflow_fit_statistical_model_to_patient.py (1)
121-131: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTighten the changed constructor typing.
The changed signature still lacks
-> None, useslist[pv.DataSet] | None, and storeslabelmap_interior_object_idsas barelist. This weakens strict mypy coverage for the new labelmap API.Proposed fix
- patient_models: list[pv.DataSet] | None = None, + patient_models: Optional[list[pv.DataSet]] = None, patient_image: Optional[itk.Image] = None, patient_labelmap: Optional[itk.Image] = None, template_labelmap: Optional[itk.Image] = None, - labelmap_interior_object_ids: Optional[list] = None, + labelmap_interior_object_ids: Optional[list[int]] = None, segmentation_method: str = "HeartSimplewareTrimmedBranches", log_level: int | str = logging.INFO, - ): + ) -> None: ... - self.labelmap_interior_object_ids: Optional[list] = labelmap_interior_object_ids + self.labelmap_interior_object_ids: Optional[list[int]] = ( + labelmap_interior_object_ids + )As per coding guidelines, use full type hints with mypy strict mode and
Optional[X]instead ofX | None.Also applies to: 168-168
🤖 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/physiomotion4d/workflow_fit_statistical_model_to_patient.py` around lines 121 - 131, The constructor typing in WorkflowFitStatisticalModelToPatient is still too loose: add an explicit return annotation to __init__, replace the PEP 604 nullable form with Optional in the patient_models parameter, and give labelmap_interior_object_ids a precise typed Optional list instead of a bare list. Update the __init__ signature and any related assignments in this class so the new labelmap API remains fully covered under strict mypy.Source: Coding guidelines
src/physiomotion4d/register_models_distance_maps.py (1)
336-350: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the disabled ICON mask code path.
Line 336 leaves the old mask transform and ICON mask calls commented out while this path now intentionally registers distance maps only. Please delete the commented-out mask code and update the nearby wording to avoid implying this is still mask-based.
♻️ Proposed cleanup
- # Pre-align moving distance map and binary mask into the fixed grid using the Greedy affine result + # Pre-align the moving distance map into the fixed grid using the + # Greedy affine result. moving_distance_map_affine_transformed = ( self.transform_tools.transform_image( self.moving_distance_map_image, @@ - # moving_mask_affine_transformed = self.transform_tools.transform_image( - # self.moving_mask_image, - # forward_transform_Greedy, - # self.reference_image, - # interpolation_method="nearest", - # ) - # Configure and run ICON self.registrar_ICON.set_number_of_iterations(icon_iterations) self.registrar_ICON.set_fixed_image(self.fixed_distance_map_image) - # self.registrar_ICON.set_fixed_mask(self.fixed_mask_image) result_ICON = self.registrar_ICON.register( moving_image=moving_distance_map_affine_transformed, - # moving_mask=moving_mask_affine_transformed, )🤖 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/physiomotion4d/register_models_distance_maps.py` around lines 336 - 350, The ICON registration path in register_models_distance_maps is now distance-map only, so remove the commented-out mask transform and mask registration lines around self.transform_tools.transform_image and self.registrar_ICON.register, and update the nearby comment/wording in this block so it no longer implies mask-based processing. Keep the active flow centered on self.registrar_ICON.set_fixed_image and register with moving_distance_map_affine_transformed only.
🤖 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 @.github/workflows/nightly-health.yml:
- Line 61: The workflow step uses a hardcoded user-specific Python executable
path, which ties nightly-health to one Windows profile on the runner. Update the
PowerShell command in the workflow to use a machine-wide interpreter or a
runner-provided variable instead of the fixed path, so the venv creation step
remains portable across self-hosted runner rebuilds and account changes. Keep
the change localized to the interpreter invocation in the affected workflow job.
In `@src/physiomotion4d/register_models_pca.py`:
- Line 140: Keep the distance-map behavior consistent between model construction
and later reassignment. Update set_fixed_model() in register_models_pca.py to
use the same squared_distance setting as the constructor path used by the
fixed-model initialization, so replacing the fixed model after construction
optimizes the same metric semantics. Use the set_fixed_model() method and the
fixed-model distance-map creation call as the main points to update.
- Around line 514-520: The rejection warning in the invalid-points check exceeds
the 88-character limit; update the warning inside the
_fixed_distance_map_max_distance path to wrap the message across multiple lines
so it stays within style limits. Keep the message construction lazy by avoiding
inline interpolation in the call to self.log_warning, and use the existing
n_invalid_points/n_valid_points values to preserve the same rejection behavior.
In `@src/physiomotion4d/workflow_fit_statistical_model_to_patient.py`:
- Around line 126-128: `WorkflowFitStatisticalModelToPatient.__init__` stores
`template_labelmap`, but `set_use_labelmap_to_image_registration` still requires
the caller to pass it again when enabling L2I. Update the setter to fall back to
the instance’s stored `template_labelmap` when the argument is omitted, and only
raise if neither the method argument nor the constructor-provided value is
available. Keep the behavior consistent in the related validation path so
`template_labelmap` is honored across the class.
- Around line 727-748: The moving mask passed into Greedy/ICON may be in a
different grid than the rebuilt template_labelmap, especially when
self.template_mask is reused after propagated_labelmap is copied into
template_labelmap geometry. Update the workflow around the template_mask
selection in workflow_fit_statistical_model_to_patient so the mask is generated
from template_labelmap or explicitly validated/resampled to match its geometry
before registrar_Greedy.register is called. Apply the same fix in the
corresponding code path for the later register call referenced by the comment.
---
Outside diff comments:
In `@src/physiomotion4d/contour_tools.py`:
- Around line 261-272: The labelmap in contour_tools.py is now using np.uint16,
but the assignment in the mesh loop can still overflow when i + 1 exceeds the
available label range. Add a bounds check in the labelmap-building logic around
create_mask_from_mesh/labelmap_arr assignment so the code fails fast or skips
with a clear error before writing an out-of-range label, keeping anatomy group
IDs consistent.
---
Nitpick comments:
In `@src/physiomotion4d/register_models_distance_maps.py`:
- Around line 336-350: The ICON registration path in
register_models_distance_maps is now distance-map only, so remove the
commented-out mask transform and mask registration lines around
self.transform_tools.transform_image and self.registrar_ICON.register, and
update the nearby comment/wording in this block so it no longer implies
mask-based processing. Keep the active flow centered on
self.registrar_ICON.set_fixed_image and register with
moving_distance_map_affine_transformed only.
In `@src/physiomotion4d/workflow_fit_statistical_model_to_patient.py`:
- Around line 121-131: The constructor typing in
WorkflowFitStatisticalModelToPatient is still too loose: add an explicit return
annotation to __init__, replace the PEP 604 nullable form with Optional in the
patient_models parameter, and give labelmap_interior_object_ids a precise typed
Optional list instead of a bare list. Update the __init__ signature and any
related assignments in this class so the new labelmap API remains fully covered
under strict mypy.
🪄 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: 1c4d1f89-ba63-41b7-9141-53e20c5bc6cc
📒 Files selected for processing (10)
.github/workflows/nightly-health.ymlexperiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.pysrc/physiomotion4d/contour_tools.pysrc/physiomotion4d/image_tools.pysrc/physiomotion4d/register_images_greedy.pysrc/physiomotion4d/register_models_distance_maps.pysrc/physiomotion4d/register_models_pca.pysrc/physiomotion4d/transform_tools.pysrc/physiomotion4d/workflow_fit_statistical_model_to_patient.pytests/test_workflow_fit_statistical_model_to_patient.py
| # nvidia-physicsnemo, which requires Python >= 3.11. | ||
| run: | | ||
| & "C:\Program Files\Python311\python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" | ||
| & "C:\Users\saylward\AppData\Local\Programs\Python\Python311/python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Avoid hardcoding a user-specific Python path.
C:\Users\saylward\AppData\Local\Programs\Python\Python311\python.exe couples CI to one developer's Windows profile on the self-hosted runner. It will break silently if that account/install is changed or the runner is rebuilt. Prefer a machine-wide interpreter (e.g. py -3.11) or a path provided via a configured runner environment variable.
♻️ Suggested change
- & "C:\Users\saylward\AppData\Local\Programs\Python\Python311/python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv"
+ py -3.11 -m venv "$env:RUNNER_TEMP\physiomotion4d-venv"📝 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.
| & "C:\Users\saylward\AppData\Local\Programs\Python\Python311/python.exe" -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" | |
| py -3.11 -m venv "$env:RUNNER_TEMP\physiomotion4d-venv" |
🤖 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 @.github/workflows/nightly-health.yml at line 61, The workflow step uses a
hardcoded user-specific Python executable path, which ties nightly-health to one
Windows profile on the runner. Update the PowerShell command in the workflow to
use a machine-wide interpreter or a runner-provided variable instead of the
fixed path, so the venv creation step remains portable across self-hosted runner
rebuilds and account changes. Keep the change localized to the interpreter
invocation in the affected workflow job.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/physiomotion4d/workflow_fit_statistical_model_to_patient.py (2)
714-728: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not leave unclassified template labels active in the L2I moving image.
Only configured foreground IDs are normalized to
1; any label not listed as background/organ/extra remains non-zero and is included byconvert_labelmap_to_mask, so unexpected labels can influence Greedy/ICON. Build the binary labelmap from the configured foreground IDs, or fail on unmapped labels.Proposed fix
- template_labelmap_arr = np.where( - np.isin(template_labelmap_arr, self.template_labelmap_background_ids), - 0, - template_labelmap_arr, - ) - template_labelmap_arr = np.where( - np.isin(template_labelmap_arr, self.template_labelmap_organ_mesh_ids), - 1, - template_labelmap_arr, - ) - template_labelmap_arr = np.where( - np.isin(template_labelmap_arr, self.template_labelmap_organ_extra_ids), - 1, - template_labelmap_arr, - ) + foreground_mask = np.isin( + template_labelmap_arr, + self.template_labelmap_organ_mesh_ids + + self.template_labelmap_organ_extra_ids, + ) + template_labelmap_arr = foreground_mask.astype(np.uint16)🤖 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/physiomotion4d/workflow_fit_statistical_model_to_patient.py` around lines 714 - 728, The template label normalization in workflow_fit_statistical_model_to_patient leaves any unmapped non-background labels active, which can leak into the L2I moving image through convert_labelmap_to_mask. Update the labelmap handling so only configured foreground IDs from template_labelmap_organ_mesh_ids and template_labelmap_organ_extra_ids are set to 1 and everything else is forced to 0, or explicitly validate and fail on any label not covered by template_labelmap_background_ids/foreground sets before proceeding.
126-143: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFully type and document the new labelmap inputs.
labelmap_interior_object_idsshould beOptional[list[int]], and the public constructor docstring should coverpatient_labelmapandtemplate_labelmap, including expected geometry.As per coding guidelines, "Use full type hints with
mypystrict mode (disallow_untyped_defs = true)" and "Write in-code docstrings instead of creating separate documentation files".Proposed fix
- labelmap_interior_object_ids: Optional[list] = None, + labelmap_interior_object_ids: Optional[list[int]] = None, @@ 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). + patient_labelmap: Optional patient labelmap in patient image physical + space. Used for labelmap-derived PCA distance maps. + template_labelmap: Optional template labelmap in template model physical + space. Used for labelmap propagation and labelmap-to-image refinement. @@ - self.labelmap_interior_object_ids: Optional[list] = labelmap_interior_object_ids + self.labelmap_interior_object_ids: Optional[list[int]] = ( + labelmap_interior_object_ids + )Also applies to: 163-168
🤖 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/physiomotion4d/workflow_fit_statistical_model_to_patient.py` around lines 126 - 143, The constructor in workflow_fit_statistical_model_to_patient should fully type the new labelmap-related inputs and document them in the class/function docstring. Update labelmap_interior_object_ids to Optional[list[int]] and add docstring entries for patient_labelmap and template_labelmap, including what geometry/role each labelmap is expected to have. Use the existing constructor and its docstring near the workflow initialization signature to keep the public API documented inline.Source: Coding guidelines
🤖 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.
Outside diff comments:
In `@src/physiomotion4d/workflow_fit_statistical_model_to_patient.py`:
- Around line 714-728: The template label normalization in
workflow_fit_statistical_model_to_patient leaves any unmapped non-background
labels active, which can leak into the L2I moving image through
convert_labelmap_to_mask. Update the labelmap handling so only configured
foreground IDs from template_labelmap_organ_mesh_ids and
template_labelmap_organ_extra_ids are set to 1 and everything else is forced to
0, or explicitly validate and fail on any label not covered by
template_labelmap_background_ids/foreground sets before proceeding.
- Around line 126-143: The constructor in
workflow_fit_statistical_model_to_patient should fully type the new
labelmap-related inputs and document them in the class/function docstring.
Update labelmap_interior_object_ids to Optional[list[int]] and add docstring
entries for patient_labelmap and template_labelmap, including what geometry/role
each labelmap is expected to have. Use the existing constructor and its
docstring near the workflow initialization signature to keep the public API
documented inline.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0d4b2be9-05ca-4ecc-b909-85e1b7ae40ba
📒 Files selected for processing (2)
src/physiomotion4d/register_models_pca.pysrc/physiomotion4d/workflow_fit_statistical_model_to_patient.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/physiomotion4d/register_models_pca.py
Summary by CodeRabbit