ENH: Add labelmap support to image registration pipeline#35
Merged
aylward merged 2 commits intoProject-MONAI:mainfrom Mar 31, 2026
Merged
ENH: Add labelmap support to image registration pipeline#35aylward merged 2 commits intoProject-MONAI:mainfrom
aylward merged 2 commits intoProject-MONAI:mainfrom
Conversation
Threads optional multi-label segmentation (labelmap) through the base class, ANTs, Greedy, ICON, and time-series registrars. When a labelmap is provided it is used in place of the binary mask; ICON activates dice-loss weighting accordingly.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughAdded optional Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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
Adds optional multi-label segmentation (labelmap) plumbing through the image registration APIs so callers can pass per-frame labelmaps (time-series) and enable labelmap-aware behavior (notably for ICON).
Changes:
- Extend base/time-series registration interfaces to accept
moving_labelmap(s)and storefixed_labelmap. - Thread
fixed_labelmapinto the ANTs/ICON registrars from the time-series pipeline and passmoving_labelmapthrough registration calls. - Update ICON registrar to prefer labelmaps over binary masks and to enable dice-loss weighting; update
docs/API_MAP.mdfor the new public API signatures.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
src/physiomotion4d/register_time_series_images.py |
Adds moving_labelmaps parameter and forwards fixed/moving labelmaps into underlying registrars. |
src/physiomotion4d/register_images_icon.py |
Adds moving_labelmap support, labelmap-as-mask behavior, and dice-loss weighting toggle. |
src/physiomotion4d/register_images_greedy.py |
Extends API signature with moving_labelmap (currently not used in implementation). |
src/physiomotion4d/register_images_base.py |
Adds stored fixed_labelmap / moving_labelmap and setter, and threads labelmap into register()/registration_method() calls. |
src/physiomotion4d/register_images_ants.py |
Extends API signature with moving_labelmap (currently not used in implementation). |
docs/API_MAP.md |
Regenerated API map to reflect the updated public method signatures. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
src/physiomotion4d/register_images_greedy.py (1)
260-267:moving_labelmapparameter added but unused in Greedy backend.Same as ANTs, the parameter is added for API consistency but is not utilized. The Greedy registration continues to use only
moving_maskfor-gm/-mmoptions. Consider adding a docstring note about this for clarity.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/physiomotion4d/register_images_greedy.py` around lines 260 - 267, In registration_method, the new parameter moving_labelmap is currently unused by the Greedy backend; update the docstring for the method (registration_method) to note that moving_labelmap is accepted for API consistency with the ANTs backend but is ignored by the Greedy implementation and that Greedy only uses moving_mask for its -gm/-mm options; reference the parameter name moving_labelmap and the function registration_method so reviewers can find and read the added clarification.src/physiomotion4d/register_images_ants.py (1)
510-517:moving_labelmapparameter added but unused.The parameter is added for API consistency with the base class, but it's not utilized in the ANTs registration. This is acceptable since ANTs doesn't natively support dice-loss weighting like ICON. However, consider:
- Adding a brief note in the docstring mentioning the parameter is accepted for API compatibility but not currently used by ANTs.
- Alternatively, logging a debug message when a labelmap is provided but ignored.
📝 Suggested docstring addition (around line 526-528)
Args: moving_image (itk.image): The 3D image to be registered/aligned. moving_mask (itk.image, optional): Binary mask defining the region of interest in the moving image + moving_labelmap (itk.image, optional): Multi-label segmentation for + the moving image. Note: Currently not used by ANTs registration; + accepted for API compatibility with other backends. moving_image_pre (ants.core.ANTsImage, optional): Pre-processed moving image🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/physiomotion4d/register_images_ants.py` around lines 510 - 517, The registration_method currently accepts moving_labelmap but never uses it; update the registration_method's docstring to note that moving_labelmap is accepted for API compatibility with the base class but is ignored by the ANTs implementation, and/or add a debug log in registration_method (using the existing logger) to emit a message when moving_labelmap is provided and will be ignored; reference the registration_method function and the moving_labelmap parameter so reviewers can find where to add the docstring note and the optional logger.debug call.src/physiomotion4d/register_images_base.py (1)
194-205: Consider validatingfixed_imagebefore settingfixed_labelmap.The
set_fixed_maskmethod validates thatfixed_imageis set before allowing a mask (line 180-181). For consistency and to catch user errors early,set_fixed_labelmapshould likely have the same guard. Additionally, consider whether setting a labelmap should clear thefixed_mask(or vice versa) to prevent inconsistent state.💡 Suggested validation
def set_fixed_labelmap(self, fixed_labelmap: Optional[itk.Image]) -> None: """Set the fixed image labelmap (multi-label segmentation). Args: fixed_labelmap (itk.Image, optional): Multi-label segmentation co-registered with the fixed image, or None to clear. """ + if fixed_labelmap is not None and self.fixed_image is None: + raise ValueError("Fixed image must be set before setting a fixed labelmap.") self.fixed_labelmap = fixed_labelmap self.forward_transform = None self.inverse_transform = None self.loss = None self.moving_image_registered = None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/physiomotion4d/register_images_base.py` around lines 194 - 205, The set_fixed_labelmap method currently sets fixed_labelmap without validating that fixed_image exists; update set_fixed_labelmap to mirror set_fixed_mask by raising an error (or ValueError) if self.fixed_image is None, and then set self.fixed_labelmap; also clear any conflicting state by resetting self.fixed_mask to None (or vice versa per project convention) along with forward_transform, inverse_transform, loss, and moving_image_registered so the object cannot be left in an inconsistent state; reference method names set_fixed_labelmap and set_fixed_mask and attributes fixed_image, fixed_labelmap, fixed_mask, forward_transform, inverse_transform, loss, moving_image_registered when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/physiomotion4d/register_images_icon.py`:
- Around line 232-248: The network is cached in self.net but dice_loss_weight is
computed per call from moving_labelmap, so subsequent calls may use the wrong
loss configuration; update register() to detect when the desired
dice_loss_weight (computed from moving_labelmap) differs from the network's
current setting and reinitialize the network if so. Concretely, store the
current dice weight on the instance (e.g., self._dice_loss_weight) when creating
the network via get_multigradicon/get_unigradicon, and on each register() call
compute dice_loss_weight from moving_labelmap and if self.net is None or
dice_loss_weight != self._dice_loss_weight then recreate the net with the new
dice_loss_weight; ensure both branches that call get_multigradicon and
get_unigradicon set self._dice_loss_weight accordingly.
- Around line 224-230: The current logic sets fixed_effective_mask based on
moving_labelmap presence which can produce mismatched masks; change the
selection so fixed_effective_mask uses self.fixed_labelmap when it exists (i.e.,
fixed_effective_mask = self.fixed_labelmap if self.fixed_labelmap is not None
else self.fixed_mask) and add a consistency check in the surrounding code (in
the same block where moving_effective_mask and fixed_effective_mask are
computed) to either raise an error or handle the case when one side has a
labelmap and the other does not so registration does not silently fall back to
unmasked behavior; reference variables: moving_effective_mask,
fixed_effective_mask, moving_labelmap, self.fixed_labelmap, moving_mask,
self.fixed_mask.
---
Nitpick comments:
In `@src/physiomotion4d/register_images_ants.py`:
- Around line 510-517: The registration_method currently accepts moving_labelmap
but never uses it; update the registration_method's docstring to note that
moving_labelmap is accepted for API compatibility with the base class but is
ignored by the ANTs implementation, and/or add a debug log in
registration_method (using the existing logger) to emit a message when
moving_labelmap is provided and will be ignored; reference the
registration_method function and the moving_labelmap parameter so reviewers can
find where to add the docstring note and the optional logger.debug call.
In `@src/physiomotion4d/register_images_base.py`:
- Around line 194-205: The set_fixed_labelmap method currently sets
fixed_labelmap without validating that fixed_image exists; update
set_fixed_labelmap to mirror set_fixed_mask by raising an error (or ValueError)
if self.fixed_image is None, and then set self.fixed_labelmap; also clear any
conflicting state by resetting self.fixed_mask to None (or vice versa per
project convention) along with forward_transform, inverse_transform, loss, and
moving_image_registered so the object cannot be left in an inconsistent state;
reference method names set_fixed_labelmap and set_fixed_mask and attributes
fixed_image, fixed_labelmap, fixed_mask, forward_transform, inverse_transform,
loss, moving_image_registered when making the change.
In `@src/physiomotion4d/register_images_greedy.py`:
- Around line 260-267: In registration_method, the new parameter moving_labelmap
is currently unused by the Greedy backend; update the docstring for the method
(registration_method) to note that moving_labelmap is accepted for API
consistency with the ANTs backend but is ignored by the Greedy implementation
and that Greedy only uses moving_mask for its -gm/-mm options; reference the
parameter name moving_labelmap and the function registration_method so reviewers
can find and read the added clarification.
🪄 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: 8cc2bdb2-809d-4c4b-9ce9-7208d5da01de
📒 Files selected for processing (6)
docs/API_MAP.mdsrc/physiomotion4d/register_images_ants.pysrc/physiomotion4d/register_images_base.pysrc/physiomotion4d/register_images_greedy.pysrc/physiomotion4d/register_images_icon.pysrc/physiomotion4d/register_time_series_images.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #35 +/- ##
=======================================
Coverage ? 13.67%
=======================================
Files ? 49
Lines ? 6596
Branches ? 0
=======================================
Hits ? 902
Misses ? 5694
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Threads optional multi-label segmentation (labelmap) through the base class, ANTs, Greedy, ICON, and time-series registrars. When a labelmap is provided it is used in place of the binary mask; ICON activates dice-loss weighting accordingly.
Summary by CodeRabbit
New Features
Chores
Documentation