rename parameter

Author: eclipse0922

monai/transforms/post/array.py

@@ -757,7 +757,7 @@ class GenerateHeatmap(Transform):
     Args:
         sigma: gaussian standard deviation. A single value is broadcast across all spatial dimensions.
         spatial_shape: optional fallback spatial shape. If ``None`` it must be provided when calling the transform.
-        truncate: extent, in multiples of ``sigma``, used to crop the gaussian support window.
+        truncated: extent, in multiples of ``sigma``, used to crop the gaussian support window.
         normalize: normalize every heatmap channel to ``[0, 1]`` when ``True``.
         dtype: target dtype for the generated heatmaps (accepts numpy or torch dtypes).
 
@@ -772,7 +772,7 @@ def __init__(
         self,
         sigma: Sequence[float] | float = 5.0,
         spatial_shape: Sequence[int] | None = None,
-        truncate: float = 3.0,
+        truncated: float = 4.0,
         normalize: bool = True,
         dtype: np.dtype | torch.dtype | type = np.float32,
     ) -> None:
@@ -784,9 +784,9 @@ def __init__(
             if float(sigma) <= 0:
                 raise ValueError("sigma must be positive.")
             self._sigma = float(sigma)
-        if truncate <= 0:
-            raise ValueError("truncate must be positive.")
-        self.truncate = float(truncate)
+        if truncated <= 0:
+            raise ValueError("truncated must be positive.")
+        self.truncated = float(truncated)
         self.normalize = normalize
         self.torch_dtype = get_equivalent_dtype(dtype, torch.Tensor)
         self.numpy_dtype = get_equivalent_dtype(dtype, np.ndarray)
@@ -816,7 +816,7 @@ def __call__(
 
         target_shape = self._resolve_spatial_shape(spatial_shape, spatial_dims)
         sigma = self._resolve_sigma(spatial_dims)
-        radius = tuple(int(np.ceil(self.truncate * s)) for s in sigma)
+        radius = tuple(int(np.ceil(self.truncated * s)) for s in sigma)
 
         heatmap = torch.zeros((batch_size, num_points, *target_shape), dtype=self.torch_dtype, device=device)
         image_bounds = tuple(int(s) for s in target_shape)

monai/transforms/post/dictionary.py

@@ -528,7 +528,7 @@ def __init__(
         heatmap_keys: KeysCollection | None = None,
         ref_image_keys: KeysCollection | None = None,
         spatial_shape: Sequence[int] | Sequence[Sequence[int]] | None = None,
-        truncate: float = 3.0,
+        truncated: float = 4.0,
         normalize: bool = True,
         dtype: np.dtype | type = np.float32,
         allow_missing_keys: bool = False,
@@ -540,7 +540,7 @@ def __init__(
         self.generator = GenerateHeatmap(
             sigma=sigma,
             spatial_shape=None,
-            truncate=truncate,
+            truncated=truncated,
             normalize=normalize,
             dtype=dtype,
         )
@@ -632,11 +632,25 @@ def _shape_from_reference(self, reference: Any, spatial_dims: int) -> tuple[int,
 
     def _prepare_output(self, heatmap: NdarrayOrTensor, reference: Any) -> Any:
         if isinstance(reference, MetaTensor):
-            converted, _, _ = convert_to_dst_type(heatmap, reference, dtype=reference.dtype, device=reference.device)
-            converted.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[1:])
+            # Use heatmap's dtype (from generator), not reference's dtype
+            converted, _, _ = convert_to_dst_type(heatmap, reference, dtype=heatmap.dtype, device=reference.device)
+            # For batched data shape is (B, C, *spatial), for non-batched it's (C, *spatial)
+            if heatmap.ndim == 5:  # 3D batched: (B, C, H, W, D)
+                converted.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[2:])
+            elif heatmap.ndim == 4:  # 2D batched (B, C, H, W) or 3D non-batched (C, H, W, D)
+                # Need to check if this is batched 2D or non-batched 3D
+                if len(heatmap.shape[1:]) == len(reference.meta.get("spatial_shape", [])):
+                    # Non-batched 3D
+                    converted.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[1:])
+                else:
+                    # Batched 2D
+                    converted.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[2:])
+            else:  # 2D non-batched: (C, H, W)
+                converted.meta["spatial_shape"] = tuple(int(v) for v in heatmap.shape[1:])
             return converted
         if isinstance(reference, torch.Tensor):
-            converted, _, _ = convert_to_dst_type(heatmap, reference, dtype=reference.dtype, device=reference.device)
+            # Use heatmap's dtype (from generator), not reference's dtype
+            converted, _, _ = convert_to_dst_type(heatmap, reference, dtype=heatmap.dtype, device=reference.device)
             return converted
         return heatmap
 

tests/transforms/test_generate_heatmap.py

@@ -223,6 +223,25 @@ def test_dict_batched_with_ref(self):
         max_vals = heatmap.max(dim=2)[0].max(dim=2)[0].max(dim=2)[0]
         np.testing.assert_allclose(max_vals.cpu().numpy(), np.ones((2, 1)), rtol=1e-5, atol=1e-5)
 
+    def test_truncated_parameter(self):
+        # Test that truncated parameter correctly controls window size
+        pt = np.array([[8.0, 8.0]], dtype=np.float32)
+        sigma = 2.0
+        
+        # Test with different truncated values
+        small_truncated = GenerateHeatmap(sigma=sigma, spatial_shape=(32, 32), truncated=2.0)(pt)[0]
+        default_truncated = GenerateHeatmap(sigma=sigma, spatial_shape=(32, 32), truncated=4.0)(pt)[0]  # default
+        large_truncated = GenerateHeatmap(sigma=sigma, spatial_shape=(32, 32), truncated=6.0)(pt)[0]
+        
+        # Larger truncated should capture more of the gaussian, resulting in slightly higher total sum
+        self.assertLess(small_truncated.sum(), default_truncated.sum())
+        self.assertLess(default_truncated.sum(), large_truncated.sum())
+        
+        # All should have same peak value (normalized to 1.0)
+        np.testing.assert_allclose(small_truncated.max(), 1.0, rtol=1e-5)
+        np.testing.assert_allclose(default_truncated.max(), 1.0, rtol=1e-5)
+        np.testing.assert_allclose(large_truncated.max(), 1.0, rtol=1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()