Add affine-aware landmark heatmap generation#8957
Conversation
|
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 (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthrough
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 (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
| def _get_reference_affine(self, reference: Any) -> torch.Tensor: | ||
| if reference is None: | ||
| raise ValueError("coordinate_space='world' requires ref_image_keys or a reference affine.") | ||
| affine = getattr(reference, "affine", None) | ||
| if affine is not None: | ||
| return affine | ||
| if isinstance(reference, (torch.Tensor, np.ndarray)) and reference.shape in ((3, 3), (4, 4)): | ||
| return reference | ||
| raise ValueError("coordinate_space='world' requires reference data with an affine matrix.") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not infer affines from raw 3x3/4x4 images.
A normal 2D reference image with shape (3, 3) or (4, 4) will be treated as an affine and inverted, causing wrong coordinates or a singular-matrix crash instead of the documented missing-affine error.
Safer fallback
affine = getattr(reference, "affine", None)
if affine is not None:
return affine
- if isinstance(reference, (torch.Tensor, np.ndarray)) and reference.shape in ((3, 3), (4, 4)):
- return reference
raise ValueError("coordinate_space='world' requires reference data with an affine matrix.")📝 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.
| def _get_reference_affine(self, reference: Any) -> torch.Tensor: | |
| if reference is None: | |
| raise ValueError("coordinate_space='world' requires ref_image_keys or a reference affine.") | |
| affine = getattr(reference, "affine", None) | |
| if affine is not None: | |
| return affine | |
| if isinstance(reference, (torch.Tensor, np.ndarray)) and reference.shape in ((3, 3), (4, 4)): | |
| return reference | |
| raise ValueError("coordinate_space='world' requires reference data with an affine matrix.") | |
| def _get_reference_affine(self, reference: Any) -> torch.Tensor: | |
| if reference is None: | |
| raise ValueError("coordinate_space='world' requires ref_image_keys or a reference affine.") | |
| affine = getattr(reference, "affine", None) | |
| if affine is not None: | |
| return affine | |
| raise ValueError("coordinate_space='world' requires reference data with an affine matrix.") |
🤖 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 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.
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
| _ERR_COORDINATE_SPACE_LEN = ( | ||
| "Argument `coordinate_space` length must match keys length when providing per-key values." | ||
| ) | ||
| _SUPPORTED_COORDINATE_SPACES = {"voxel", "world"} |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Make the coordinate-space set immutable.
Ruff flags this mutable class attribute. Use frozenset.
Proposed fix
- _SUPPORTED_COORDINATE_SPACES = {"voxel", "world"}
+ _SUPPORTED_COORDINATE_SPACES = frozenset({"voxel", "world"})📝 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.
| _SUPPORTED_COORDINATE_SPACES = {"voxel", "world"} | |
| _SUPPORTED_COORDINATE_SPACES = frozenset({"voxel", "world"}) |
🧰 Tools
🪛 Ruff (0.15.18)
[warning] 616-616: Mutable default value for class attribute
(RUF012)
🤖 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 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.
Source: Linters/SAST tools
Signed-off-by: Mustafa Merchant <mustafamerchant072@gmail.com>
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.