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 (1)
🚧 Files skipped from review as they are similar to previous changes (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
| 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.
There was a problem hiding this comment.
I think this is a valid concern. If you want reference to possibly be the affine itself you will need some other mechanism to indicate this.
There was a problem hiding this comment.
@ericspod Agreed. The current shape-based inference is ambiguous; direct affine input needs an explicit API. Removing that fallback is the safe fix for this PR.
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.
| - Nested dot-notation key access in `ConfigParser`. | ||
| - Auto3DSeg algo serialization migrated from pickle to JSON for improved security and portability. | ||
| - Global coordinates support in spatial crop transforms. These now support global coordinate mode, allowing crops to be specified in world/global coordinates rather than local image indices, improving interoperability with physical-space annotations. | ||
| - `GenerateHeatmapd` can convert world-coordinate landmarks to reference-image voxel space and emit landmark visibility masks. |
There was a problem hiding this comment.
Version 1.6 has been released so this shouldn't be included here.
| sigma: Sequence[float] | float = 5.0, | ||
| heatmap_keys: KeysCollection | None = None, | ||
| ref_image_keys: KeysCollection | None = None, | ||
| coordinate_space: str | Sequence[str] = "voxel", |
There was a problem hiding this comment.
New arguments should be at the end of the argument list so that any existing invocation of this constructor using positional arguments isn't broken. The order of the docstrings should match as well.
| 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.
I think this is a valid concern. If you want reference to possibly be the affine itself you will need some other mechanism to indicate this.
| 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.
| heatmap.affine = reference.affine | ||
| self._update_spatial_metadata(heatmap, shape) | ||
| d[out_key] = heatmap | ||
| if visibility_key is not None: |
There was a problem hiding this comment.
If visibility_key is None then the computation for visiblity is done needlessly, can line 652 be placed in this condition block?
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.