Class GenerateHeatmap

monai/transforms/post/array.py:753–898

Generate per-landmark Gaussian heatmaps for 2D or 3D coordinates. Notes: - Coordinates are interpreted in voxel units and expected in (Y, X) for 2D or (Z, Y, X) for 3D. - Target spatial_shape is (Y, X) for 2D and (Z, Y, X) for 3D. - Output layout uses channel-first

class GenerateHeatmap(Transform):
    """
    Generate per-landmark Gaussian heatmaps for 2D or 3D coordinates.

    Notes:
        - Coordinates are interpreted in voxel units and expected in (Y, X) for 2D or (Z, Y, X) for 3D.
        - Target spatial_shape is (Y, X) for 2D and (Z, Y, X) for 3D.
        - Output layout uses channel-first convention with one channel per landmark.
        - Input points shape: (N, D) where N is number of landmarks, D is spatial dimensions (2 or 3).
        - Output heatmap shape: (N, Y, X) for 2D or (N, Z, Y, X) for 3D.
        - Each channel index corresponds to one landmark.

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

    Raises:
        ValueError: when ``sigma`` is non-positive or ``spatial_shape`` cannot be resolved.

    """

    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]

    def __init__(
        self,
        sigma: Sequence[float] | float = 5.0,
        spatial_shape: Sequence[int] | None = None,
        truncated: float = 4.0,
        normalize: bool = True,
        dtype: np.dtype | torch.dtype | type = np.float32,
    ) -> None:
        if isinstance(sigma, Sequence) and not isinstance(sigma, (str, bytes)):
            if any(s <= 0 for s in sigma):
                raise ValueError("Argument `sigma` values must be positive.")
            self._sigma = tuple(float(s) for s in sigma)
        else:
            if float(sigma) <= 0:
                raise ValueError("Argument `sigma` must be positive.")
            self._sigma = (float(sigma),)
        if truncated <= 0:
            raise ValueError("Argument `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)
        # Validate that dtype is floating-point for meaningful Gaussian values
        if not self.torch_dtype.is_floating_point:
            raise ValueError(f"Argument `dtype` must be a floating-point type, got {self.torch_dtype}")
        self.spatial_shape = None if spatial_shape is None else tuple(int(s) for s in spatial_shape)

    def __call__(self, points: NdarrayOrTensor, spatial_shape: Sequence[int] | None = None) -> NdarrayOrTensor:
        """
        Args:
            points: landmark coordinates as ndarray/Tensor with shape (N, D),
                ordered as (Y, X) for 2D or (Z, Y, X) for 3D, where N is the number