Skip to content

Add affine-aware landmark heatmap generation#8957

Open
mustafamm072 wants to merge 2 commits into
Project-MONAI:devfrom
mustafamm072:dev
Open

Add affine-aware landmark heatmap generation#8957
mustafamm072 wants to merge 2 commits into
Project-MONAI:devfrom
mustafamm072:dev

Conversation

@mustafamm072

Copy link
Copy Markdown

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 using ref_image_keys before heatmap generation. The transform can also emit per-landmark visibility masks via visibility_keys.

Related to #7486.

Testing

  • Added tests for translated, rotated, anisotropic, and scaled reference affines.
  • Added tests for MetaTensor point affines and visibility masks.
  • Ran py_compile and git diff --check.

@coderabbitai

coderabbitai Bot commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 98cfa754-135d-4bce-a6ee-19e3ece6795e

📥 Commits

Reviewing files that changed from the base of the PR and between dbbfbb3 and 27f4be3.

📒 Files selected for processing (2)
  • monai/transforms/post/dictionary.py
  • tests/transforms/test_generate_heatmapd.py
🚧 Files skipped from review as they are similar to previous changes (2)
  • tests/transforms/test_generate_heatmapd.py
  • monai/transforms/post/dictionary.py

📝 Walkthrough

Walkthrough

GenerateHeatmapd now accepts coordinate_space and visibility_keys. It can convert world-space landmarks into voxel space using a reference affine or a MetaTensor affine on the points, then generate heatmaps and per-point visibility masks. Validation was added for unsupported coordinate spaces, missing affine data, and mismatched key lengths. Tests cover world-space conversion, visibility behavior, and error cases.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 10.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the main change: affine-aware landmark heatmap generation.
Description check ✅ Passed It covers the feature, default behavior, related issue, and testing, but it omits the template's checklist section and explicit Fixes # reference.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (2)
monai/transforms/post/dictionary.py (1)

616-616: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make the supported coordinate spaces immutable.

Ruff flags this mutable class attribute; frozenset matches 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 win

Cover the remaining new public contracts.

Please add direct tests for coordinate_space length mismatch and NumPy visibility output; both are new branches in GenerateHeatmapd.

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

📥 Commits

Reviewing files that changed from the base of the PR and between b7d14c8 and b71a614.

📒 Files selected for processing (3)
  • docs/source/whatsnew_1_6.md
  • monai/transforms/post/dictionary.py
  • tests/transforms/test_generate_heatmapd.py

Comment thread monai/transforms/post/dictionary.py
Comment on lines +796 to +804
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.")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.

Suggested change
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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
monai/transforms/post/dictionary.py (1)

618-641: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add 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

📥 Commits

Reviewing files that changed from the base of the PR and between b71a614 and dbbfbb3.

📒 Files selected for processing (3)
  • docs/source/whatsnew_1_6.md
  • monai/transforms/post/dictionary.py
  • tests/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"}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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.

Suggested change
_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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant