生成二维高斯分布热力图

背景

这篇博文是因为想复现这篇博客远距离行人检测：从标注方式出发。它使用了两个点来标注行人，并用高斯分布为两个点生成热力图，作为模型预测目标。那么很自然，第一个问题就是高斯分布热力图如何生成。其中高斯分布$e$项之前的$\frac 1 {\sqrt {2\pi} \sigma}$系数被省去，因为不同尺度的行人拟合的分布的$\sigma$必定不同，省去才能使热力图的值一致。代码写好了，数据集还没下载完，就撸篇博客吧。

自己的实现（逐点运算）

看了一些博客后自己实现了一个版本，代码如下：

%matplotlib inline
import numpy as np
from math import exp, log, sqrt, ceil
from matplotlib import pyplot as plt

def cal_sigma(dmax, edge_value):
    """Calculate a sigma makes those point's distance to 
    center is dmax value equal to edge value"""
    return sqrt(- pow(dmax, 2) / log(edge_value))

def gaussian(dx, dy, sigma):
    """modifyed version normal distribution pdf"""
    x_term = pow(dx, 2)
    y_term = pow(dy, 2)
    exp_value = - (x_term + y_term) / 2 / pow(sigma, 2)
    return exp(exp_value)

def draw_heatmap(width, height, dmax, x, y, sigma, edge_value):
    heatmap = np.zeros((height, width), dtype=np.float32)
    r = sqrt(2 * pow(dmax, 2))
    intr = int(1.9*dmax)
    xmin = max(0, x - intr)
    xmax = min(width - 1, x + intr)
    ymin = max(0, y - intr)
    ymax = min(height - 1, y + intr)
    for i in range(xmin, xmax):
        for j in range(ymin, ymax):
            dx = x - i
            dy = y - j
            dist = sqrt(pow(dx, 2) + pow(dy, 2))
            if dist > intr:
                continue
            value = gaussian(x - i, y - j, sigma)
            heatmap[j, i] = value
    return heatmap
    
def test():
    dmax = 100
    edge_value = 0.01
    sigma = cal_sigma(dmax, edge_value)
    hm = draw_heatmap(1024, 576, dmax, 512, 288, sigma, edge_value)
    return hm
    
# %timeit hm = test()
hm = test()
plt.imshow(hm)
plt.show()

分布的$\sigma$可以通过想要拟合的大小及改点的值算出，也就是代码中的$cal_sigma$函数。gaussian是修改版的二维高斯分布概率密度函数（Probability Density Function，pdf），逐点运算版本。为了提高速度，在$draw_heatmap$中对遍历范围进行了约束。

生成的热力图如下：

Numpy版本（向量计算）

逐点遍历速度当然不如向量计算来得快，通过面向Google编程后，用scipy的multivariate_normal方法实现了向量版。但因为我需要修改版的高斯分布，故用numpy实现了向量计算。代码如下：

import numpy as np
from math import exp, log, sqrt, ceil
from matplotlib import pyplot as plt

# from scipy.stats import multivariate_normal

def cal_sigma(dmax, edge_value):
    return sqrt(- pow(dmax, 2) / log(edge_value))

def gaussian(array_like_hm, mean, sigma):
    """modifyed version normal distribution pdf, vector version"""
    array_like_hm -= mean
    x_term = array_like_hm[:,0] ** 2
    y_term = array_like_hm[:,1] ** 2
    exp_value = - (x_term + y_term) / 2 / pow(sigma, 2)
    return np.exp(exp_value)
  
def draw_heatmap(width, height, x, y, sigma, array_like_hm):
    m1 = (x, y)
    s1 = np.eye(2) * pow(sigma, 2)
    # k1 = multivariate_normal(mean=m1, cov=593.109206084)
    k1 = multivariate_normal(mean=m1, cov=s1)
    #     zz = k1.pdf(array_like_hm)
    zz = gaussian(array_like_hm, m1, sigma)
    img = zz.reshape((height,width))
    return img

def test(width, height, x, y, array_like_hm):
    dmax = 100
    edge_value = 0.01
    sigma = cal_sigma(dmax, edge_value)
    
    return draw_heatmap(width, height, x, y, sigma, array_like_hm)

xres = 1024
yres = 576
xlim = (0, xres)
ylim = (0, yres)

# x = np.linspace(xlim[0], xlim[1], xres)
# y = np.linspace(ylim[0], ylim[1], yres)
x = np.arange(xres, dtype=np.float)
y = np.arange(yres, dtype=np.float)
xx, yy = np.meshgrid(x,y)

# evaluate kernels at grid points
xxyy = np.c_[xx.ravel(), yy.ravel()]
# %timeit img = test(1024, 576, 512, 288, xxyy.copy())
img = test(xres, yres, xres / 2, yres / 2, xxyy.copy())
plt.imshow(img)
plt.show()

这里的gaussion函数就是向量版本的。热力图如下：

速度测试

测试显示逐点版本在这个参数下计算一次高斯热力图需要142ms，而向量版本仅需21.4ms。逐点版本已经通过坐标限制了计算范围了，向量版本是对全图进行的计算就已领先这么多，如果也限制计算范围，估计速度还能提升一大截。

作为多年的程序员，遍历已经写了不知道多少次了。啥问题拿过来脑子里想的都是怎么遍历。然而很多时候矩阵计算才是王道，还需好好学习科学计算的各种库，学习矩阵计算的思想。

参考资料

How to efficiently compute the heat map of two Gaussian distribution in Python?
tf-openpose人体姿态估计标签生成—heatmap—vectormap