Add affine-aware landmark heatmap generation#8957
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthrough
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![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 (2)
monai/transforms/post/dictionary.py (1)
616-616: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the supported coordinate spaces immutable.
Ruff flags this mutable class attribute;
frozensetmatches constant intent.Proposed fix
- _SUPPORTED_COORDINATE_SPACES = {"voxel", "world"} + _SUPPORTED_COORDINATE_SPACES = frozenset({"voxel", "world"})🤖 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 `@monai/transforms/post/dictionary.py` at line 616, The _SUPPORTED_COORDINATE_SPACES class constant in the coordinate-space transform logic is currently a mutable set; replace it with an immutable frozenset to match constant intent and satisfy Ruff. Update the definition in the relevant class within dictionary.py where _SUPPORTED_COORDINATE_SPACES is declared, keeping the same values ("voxel" and "world") but making the attribute immutable.Source: Linters/SAST tools
tests/transforms/test_generate_heatmapd.py (1)
317-328: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the remaining new public contracts.
Please add direct tests for
coordinate_spacelength mismatch and NumPy visibility output; both are new branches inGenerateHeatmapd.Example additions
+ def test_coordinate_space_length_mismatch_raises(self): + with self.assertRaises(ValueError): + GenerateHeatmapd( + keys=["pts1", "pts2"], + heatmap_keys=["hm1", "hm2"], + coordinate_space=["voxel", "world", "voxel"], + spatial_shape=(8, 8), + ) + + def test_numpy_points_visibility_is_numpy(self): + image = MetaTensor(torch.zeros((1, 8, 8, 8), dtype=torch.float32), affine=torch.eye(4)) + image.meta["spatial_shape"] = (8, 8, 8) + points = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + + result = GenerateHeatmapd( + keys="points", + heatmap_keys="heatmap", + ref_image_keys="image", + coordinate_space="world", + visibility_keys="visible", + sigma=1.0, + )({"points": points, "image": image}) + + self.assertIsInstance(result["visible"], np.ndarray) + self.assertEqual(result["visible"].dtype, np.bool_)As per path instructions,
**/*.py: “Ensure new or modified definitions will be covered by existing or new unit tests.”🤖 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 `@tests/transforms/test_generate_heatmapd.py` around lines 317 - 328, Add direct coverage for the remaining public branches in GenerateHeatmapd: the constructor should reject a coordinate_space length mismatch, and __call__ should be tested for NumPy visibility output when visibility_keys is provided. Extend the existing tests in test_generate_heatmapd.py to exercise these cases using GenerateHeatmapd with mismatched coordinate_space alongside the current visibility_keys mismatch test, and add an assertion that the visibility output is a NumPy array for the relevant GenerateHeatmapd invocation.Source: Path instructions
🤖 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 `@monai/transforms/post/dictionary.py`:
- Around line 796-804: The _get_reference_affine helper in the dictionary
transform is incorrectly treating any raw 3x3/4x4 tensor or ndarray as an
affine, which can misclassify normal reference images. Update this logic so only
objects with an explicit affine attribute are accepted for
coordinate_space='world', and otherwise raise the documented missing-affine
ValueError instead of inferring from shape; keep the fix localized to
_get_reference_affine and its callers.
- Around line 583-591: The world-space example in GenerateHeatmapd is
inconsistent with the sample data: it references landmarks_world even though the
example data only defines landmarks, and it uses an ndarray image without affine
metadata. Update the doc example to either define the matching world-space
landmark key and a reference image object with affine information, or change the
example to use the existing landmarks with a valid image/affine setup so the
GenerateHeatmapd usage is runnable.
---
Nitpick comments:
In `@monai/transforms/post/dictionary.py`:
- Line 616: The _SUPPORTED_COORDINATE_SPACES class constant in the
coordinate-space transform logic is currently a mutable set; replace it with an
immutable frozenset to match constant intent and satisfy Ruff. Update the
definition in the relevant class within dictionary.py where
_SUPPORTED_COORDINATE_SPACES is declared, keeping the same values ("voxel" and
"world") but making the attribute immutable.
In `@tests/transforms/test_generate_heatmapd.py`:
- Around line 317-328: Add direct coverage for the remaining public branches in
GenerateHeatmapd: the constructor should reject a coordinate_space length
mismatch, and __call__ should be tested for NumPy visibility output when
visibility_keys is provided. Extend the existing tests in
test_generate_heatmapd.py to exercise these cases using GenerateHeatmapd with
mismatched coordinate_space alongside the current visibility_keys mismatch test,
and add an assertion that the visibility output is a NumPy array for the
relevant GenerateHeatmapd invocation.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: d6e314af-388d-48b3-a672-b87902ab2768
📒 Files selected for processing (3)
docs/source/whatsnew_1_6.mdmonai/transforms/post/dictionary.pytests/transforms/test_generate_heatmapd.py
Signed-off-by: Mustafa Merchant <mustafamerchant072@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
monai/transforms/post/dictionary.py (1)
618-641: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Google-style docstrings to the changed definitions.
The new/modified methods should document Args, Returns, and raised ValueErrors.
As per path instructions,
**/*.py: “Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings.”Also applies to: 643-672, 695-721, 780-818
🤖 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 `@monai/transforms/post/dictionary.py` around lines 618 - 641, Add Google-style docstrings to the changed definitions in dictionary.py, including the constructor and the related helper/transform methods referenced by the review, so each one documents Args, Returns, and any ValueError raised. Update the docstrings on the affected symbols such as __init__, _prepare_heatmap_keys, _prepare_optional_keys, _prepare_coordinate_spaces, _prepare_visibility_keys, _prepare_shapes, and any other modified definitions in the same area, keeping the descriptions aligned with the existing parameter names and behavior.Source: Path instructions
🤖 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 `@monai/transforms/post/dictionary.py`:
- Line 616: The _SUPPORTED_COORDINATE_SPACES class attribute in the
coordinate-space transform code is a mutable set and should be made immutable.
Update the attribute to use frozenset in the same class/module where it is
defined so Ruff no longer flags it, keeping the existing "voxel" and "world"
values unchanged.
---
Nitpick comments:
In `@monai/transforms/post/dictionary.py`:
- Around line 618-641: Add Google-style docstrings to the changed definitions in
dictionary.py, including the constructor and the related helper/transform
methods referenced by the review, so each one documents Args, Returns, and any
ValueError raised. Update the docstrings on the affected symbols such as
__init__, _prepare_heatmap_keys, _prepare_optional_keys,
_prepare_coordinate_spaces, _prepare_visibility_keys, _prepare_shapes, and any
other modified definitions in the same area, keeping the descriptions aligned
with the existing parameter names and 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 2ec5fb77-ebfa-4a14-a352-002f6bae4b54
📒 Files selected for processing (3)
docs/source/whatsnew_1_6.mdmonai/transforms/post/dictionary.pytests/transforms/test_generate_heatmapd.py
✅ Files skipped from review due to trivial changes (1)
- docs/source/whatsnew_1_6.md
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/transforms/test_generate_heatmapd.py
Signed-off-by: Mustafa Merchant <mustafamerchant072@gmail.com>
Signed-off-by: Mustafa Merchant <mustafamerchant072@gmail.com>
ericspod
left a comment
ericspod
left a comment
There was a problem hiding this comment.
hi @mustafamm072 thanks for this addition, I had some comments but with changes we should be good to go.
| if points_t.ndim != 2: | ||
| raise ValueError(f"{self._ERR_INVALID_POINTS} Got {points_t.ndim}D tensor.") | ||
| bounds = torch.as_tensor(spatial_shape, dtype=points_t.dtype, device=points_t.device) | ||
| return torch.isfinite(points_t).all(dim=1) & (points_t >= 0).all(dim=1) & (points_t < bounds).all(dim=1) |
There was a problem hiding this comment.
| return torch.isfinite(points_t).all(dim=1) & (points_t >= 0).all(dim=1) & (points_t < bounds).all(dim=1) | |
| return points_t.mul(torch.isfinite(points_t)).clamp_(0, bounds - 1).all(dim=1) |
My concern is that the original expression involves a lot of tensor copies that's not efficient. Please check this is correct if you think it's worth changing.
There was a problem hiding this comment.
Reduced the copies by combining the checks into a single boolean mask and doing one .all(dim=1) reduction. I didn't use the exact points_t.mul(...).clamp_(0, bounds-1) form because it changes the semantics. clamp_ forces every point in bounds so all points would report visible, and multiplying by the finite-mask doesn't drop NaNs (0.0 * nan == nan). The new version keeps correctness while removing the redundant reductions. Let me know if you'd prefer a different approach.
…oordinate support - Move coordinate_space/visibility_keys to end of __init__ to preserve positional-argument compatibility; reorder docstring to match - Remove ambiguous shape-based affine inference from raw 3x3/4x4 arrays; require an explicit .affine attribute - Compute visibility mask only when visibility_keys is provided - Build a single boolean mask in _compute_visibility to reduce tensor copies - Use frozenset for _SUPPORTED_COORDINATE_SPACES (Ruff RUF012) - Make the world-space docstring example runnable - Drop whatsnew_1_6 entry (1.6 already released) Signed-off-by: Mustafa Merchant <mustafamerchant072@gmail.com>
2 participants
Description
Adds opt-in world-coordinate landmark support to
GenerateHeatmapd.By default, behavior is unchanged. With
coordinate_space="world", landmarks are transformed into the reference image voxel space usingref_image_keysbefore heatmap generation. The transform can also emit per-landmark visibility masks viavisibility_keys.Related to #7486.
Testing
MetaTensorpoint affines and visibility masks.py_compileandgit diff --check.