MONAI
MONAI copied to clipboard
Generate heatmap transforms
Fixes #3328 .
Description
A few sentences describing the changes proposed in this pull request.
This pull request introduces GenerateHeatmap and GenerateHeatmapd transforms for creating Gaussian heatmaps from landmark coordinates.
The input points are currently expected in ZYX order, but this can be changed to support XYZ if preferred.
The transforms support both batched (B, N, D) and non-batched (N, D) inputs.
Example notebooks are included for demonstration and will be removed before the PR is merged.
Types of changes
- [x] Non-breaking change (fix or new feature that would not break existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing functionality to change).
- [x] New tests added to cover the changes.
- [x] Integration tests passed locally by running
./runtests.sh -f -u --net --coverage. - [x] Quick tests passed locally by running
./runtests.sh --quick --unittests --disttests. - [ ] In-line docstrings updated.
- [ ] Documentation updated, tested
make htmlcommand in thedocs/folder.
[!NOTE]
Other AI code review bot(s) detected
CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.
Walkthrough
Adds a new public GenerateHeatmap transform (monai/transforms/post/array.py) that creates per-landmark Gaussian heatmaps for 2D/3D batched coordinates with flexible sigma handling, optional truncation, per-channel normalization, channel-first output, and dtype mapping via get_equivalent_dtype. Adds GenerateHeatmapd (monai/transforms/post/dictionary.py) with aliases GenerateHeatmapD/GenerateHeatmapDict to apply the array transform per key, infer spatial_shape (constructor or per-call) or use a reference image, convert output dtype/device via convert_to_dst_type, and propagate MetaTensor metadata when available. Comprehensive unit tests for both variants were added and package exports (all) updated.
Estimated code review effort
🎯 4 (Complex) | ⏱️ ~45 minutes
- monai/transforms/post/array.py: verify Gaussian computation, sigma broadcasting/resolution, truncation/window logic, spatial_shape resolution, channel ordering, dtype selection/mapping, and GaussianFilter integration.
- monai/transforms/post/dictionary.py: review MapTransform wrapper behavior, per-key spatial_shape inference/validation, reference image dtype/device/MetaTensor handling, use of convert_to_dst_type, and multi-key error paths.
- tests/transforms/test_generate_heatmap.py and tests/transforms/test_generate_heatmapd.py: run and inspect tests for numerical tolerances (peak alignment, normalization), dtype/device propagation, batched and multi-key scenarios, and invalid-input assertions.
Pre-merge checks and finishing touches
❌ Failed checks (1 warning, 1 inconclusive)
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage | ⚠️ Warning | Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. | You can run @coderabbitai generate docstrings to improve docstring coverage. |
| Description check | ❓ Inconclusive | PR description matches template structure with all required sections completed; however, strikethrough text indicates incomplete implementation details. | Clarify whether batched inputs are supported (strikethrough suggests uncertainty) and confirm docstrings/documentation status before merge. |
✅ Passed checks (3 passed)
| Check name | Status | Explanation |
|---|---|---|
| Title check | ✅ Passed | Title accurately describes the main change—introducing GenerateHeatmap and GenerateHeatmapd transforms for heatmap generation. |
| Linked Issues check | ✅ Passed | PR implements the core requirements from #3328: GenerateHeatmap and GenerateHeatmapd transforms, configurable sigma/size parameters, dictionary-based out_key pattern, and support for batched/non-batched inputs. |
| Out of Scope Changes check | ✅ Passed | All changes directly support heatmap transform implementation—new transforms, tests, imports (get_equivalent_dtype, convert_to_dst_type), and public API exports. No extraneous changes detected. |
✨ Finishing touches
- [ ] 📝 Generate docstrings
🧪 Generate unit tests (beta)
- [ ] Create PR with unit tests
- [ ] Post copyable unit tests in a comment
[!TIP]
📝 Customizable high-level summaries are now available in beta!
You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
- Provide your own instructions using the
high_level_summary_instructionssetting.- Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
- Use
high_level_summary_in_walkthroughto move the summary from the description to the walkthrough section.Example instruction:
"Divide the high-level summary into five sections:
- 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
- 📓 References — List relevant issues, discussions, documentation, or related PRs.
- 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
- 📊 Contributor Summary — Include a Markdown table showing contributions:
| Contributor | Lines Added | Lines Removed | Files Changed |- ✔️ Additional Notes — Add any extra reviewer context. Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.
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 help to get the list of available commands and usage tips.
Could you please take a look at this PR when you have a moment? I believe the GenerateHeatmap would be a valuable addition. Thanks.
@KumoLiu @ericspod @Nic-Ma
Hi @eclipse0922 sorry for the delay in reviewing. I think these classes would be a good addition to MONAI, I had a few comments below but I hadn't gone through the code thoroughly quiet yet. I commented below that these classes could use GaussianFilter instead of defining the gaussian computation over again, this would save reduce line count and so complexity, and should be adaptable to how you're computing the heatmaps. Please have a look at that class, or elsewhere in MONAI that gaussian kernels are used. Thanks!
@ericspod
Thank you for the detailed review. I'll look into your comments and incorporate the changes, including the suggestion to use GaussianFilter
2d_mdtest.ipynb 3d_heatmap_pyvista.ipynb updated notebook
Hi @ericspod, just a friendly reminder regarding this PR.
I have addressed your feedback from the last review, specifically refactoring the code to use GaussianFilter instead of the custom implementation. I also fixed the issues reported by the AI review bots.
Could you please take a look again when you have a chance? Thanks!
/build