import numpy as np def draw_gauss(heatmap, x, y, gsize): height, width = heatmap.shape[:2] gsize += 1 - (gsize % 2) sigma = gsize / 6 s = 2 * sigma * sigma radius = gsize // 2 ky, kx = np.ogrid[-radius:+radius+1, -radius:+radius+1] kernel = np.exp(-(kx * kx + ky * ky) / s) dleft, dtop = -min(x, radius), -min(y, radius) dright, dbottom = +min(width - x, radius+1), +min(height - y, radius+1) select_heatmap = heatmap[y+dtop:y+dbottom, x+dleft:x+dright] select_kernel = kernel[radius+dtop:radius+dbottom, radius+dleft:radius+dright] if min(select_heatmap.shape) > 0: np.maximum(select_heatmap, select_kernel, out=select_heatmap) return heatmap