How to Generate Gaussian Heatmap on numpy

gaussian_heatmap.py

import cv2
import numpy


def gaussian_heatmap(sigma: int, spread: int):
    extent = int(spread * sigma)
    center = spread * sigma / 2
    heatmap = numpy.zeros([extent, extent], dtype=numpy.float32)
    for i_ in range(extent):
        for j_ in range(extent):
            heatmap[i_, j_] = 1 / 2 / numpy.pi / (sigma ** 2) * numpy.exp(
                -1 / 2 * ((i_ - center - 0.5) ** 2 + (j_ - center - 0.5) ** 2) / (sigma ** 2))
    heatmap = (heatmap / numpy.max(heatmap) * 255).astype(numpy.uint8)
    return heatmap


hm = gaussian_heatmap(sigma=10, spread=3)

hm = cv2.applyColorMap(hm, cv2.COLORMAP_JET)
cv2.imshow("hm", hm)
cv2.waitKey(0)