Fixes heatmap normalization and shape checking

Fixes an issue where heatmap normalization was using the entire heatmap instead of the local region.

Adds a check to ensure that the provided static shape matches the number of spatial dimensions.

Signed-off-by: sewon.jeon <[email protected]>

Author: eclipse0922

monai/transforms/post/array.py

@@ -841,9 +841,9 @@ def __call__(self, points: NdarrayOrTensor, spatial_shape: Sequence[int] | None
                 # write back
                 region.copy_(updated)
                 if self.normalize:
-                    peak = updated.amax()
+                    peak = heatmap[b_idx, idx].amax()
                     denom = torch.where(peak > 0, peak, torch.ones_like(peak))
-                    heatmap[b_idx, idx] = heatmap[b_idx, idx] / denom
+                    heatmap[b_idx, idx].div_(denom)
 
         if not is_batched:
             heatmap = heatmap.squeeze(0)

monai/transforms/post/dictionary.py

@@ -520,7 +520,8 @@ class GenerateHeatmapd(MapTransform):
 
     Args:
         keys: keys of the corresponding items in the dictionary.
-        sigma: standard deviation for the Gaussian kernel. Can be a single value or sequence matching number of points.
+        sigma: standard deviation for the Gaussian kernel. Can be a single value or a sequence matching the number
+            of spatial dimensions.
         heatmap_keys: keys to store output heatmaps. Default: "{key}_heatmap" for each key.
         ref_image_keys: keys of reference images to inherit spatial metadata from. When provided, heatmaps will
             have the same shape, affine, and spatial metadata as the reference images.
@@ -647,12 +648,16 @@ def _prepare_shapes(
     def _determine_shape(
         self, points: Any, static_shape: tuple[int, ...] | None, data: Mapping[Hashable, Any], ref_key: Hashable | None
     ) -> tuple[int, ...]:
-        if static_shape is not None:
-            return static_shape
         points_t = convert_to_tensor(points, dtype=torch.float32, track_meta=False)
         if points_t.ndim not in (2, 3):
             raise ValueError(f"{self._ERR_INVALID_POINTS} Got {points_t.ndim}D tensor.")
         spatial_dims = int(points_t.shape[-1])
+        if static_shape is not None:
+            if len(static_shape) != spatial_dims:
+                raise ValueError(
+                    f"Provided static spatial_shape has {len(static_shape)} dims; expected {spatial_dims}."
+                )
+            return static_shape
         if ref_key is not None and ref_key in data:
             return self._shape_from_reference(data[ref_key], spatial_dims)
         raise ValueError(self._ERR_NO_SHAPE)