Centernet 生成高斯热图

写在前面的话

最近学校阳了，宿舍给封了，宿舍网络不好远程跑不了实验，随缘写一下对CenterNet源码的一个解读，之前写论文的那段时间留下来的工作，respect！

这个文章主要是对CenterNet中生成高斯核的部分代码进行解析，具体原理不会细讲，但是本文增加了一个很方便理解的可视化的代码，可以自己拿来跑就行，自己debug应该也可以理解作者的意思，希望对读者有帮助。

可视化代码下载链接：https://download.csdn.net/download/weixin_42899627/87157112

Centernet 源码位置
本文核心代码在CenterNet/src/lib/utils/image.py中可以找到

二维高斯函数的公式

CenterNet源码中二维高斯函数实现如下：

tip： 对比公式少了些东西，但是不影响高斯函数的特性，这里关键还是看高斯核半径的计算

defgaussian2D(shape, sigma=1):
    m, n =[(ss -1.)/2.for ss in shape]
    y, x = np.ogrid[-m:m +1,-n:n +1]#np.orgin 生成二维网格坐标

    h = np.exp(-(x * x + y * y)/(2* sigma * sigma))
    h[h < np.finfo(h.dtype).eps * h.max()]=0#np.finfo()常用于生成一定格式，数值较小的偏置项eps，以避免分母或对数变量为零return h

高斯核半径的计算

从代码上看就是一元二次方程的求根公式

这里要注意的代码中计算高斯半径是根据框的角点进行计算，而在Centernet中需要计算的是框的中心点的高斯半径，其实道理是一样的 Centernet 框的角点的偏移可以近似对于框中心点的偏移

情况一：两角点均在真值框内
情况二：两角点均在真值框外
情况三：一角点在真值框内，一角点在真值框外

参考文章：
CornerNet Guassian radius高斯半径的确定-数学公式详解
说点Cornernet/Centernet代码里面GT heatmap里面如何应用高斯散射核

defgaussian_radius(det_size, min_overlap=0.7):
    height, width = det_size

    a1 =1
    b1 =(height + width)
    c1 = width * height *(1- min_overlap)/(1+ min_overlap)
    sq1 = np.sqrt(b1 **2-4* a1 * c1)
    r1 =(b1 + sq1)/2

    a2 =4
    b2 =2*(height + width)
    c2 =(1- min_overlap)* width * height
    sq2 = np.sqrt(b2 **2-4* a2 * c2)
    r2 =(b2 + sq2)/2

    a3 =4* min_overlap
    b3 =-2* min_overlap *(height + width)
    c3 =(min_overlap -1)* width * height
    sq3 = np.sqrt(b3 **2-4* a3 * c3)
    r3 =(b3 + sq3)/2returnmin(r1, r2, r3)

CenterNet源码中 draw_umich_gaussian 函数实现如下：

tip: 没啥特别的操作，主要是将生成的一个二维高斯核（目标框尺寸）放到原图（图像尺寸）的对应位置上

defdraw_umich_gaussian(heatmap, center, radius, k=1):
    diameter =2* radius +1
    gaussian = gaussian2D((diameter, diameter), sigma=diameter /6)

    x, y =int(center[0]),int(center[1])

    height, width = heatmap.shape[0:2]

    left, right =min(x, radius),min(width - x, radius +1)
    top, bottom =min(y, radius),min(height - y, radius +1)

    masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]ifmin(masked_gaussian.shape)>0andmin(masked_heatmap.shape)>0:# TODO debug
        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)#逐个元素比较大小，保留大的值return heatmap

import numpy as np
import math
import xml.etree.ElementTree as ET
import glob
from image import draw_dense_reg, draw_msra_gaussian, draw_umich_gaussian
from image import get_affine_transform, affine_transform, gaussian_radius

data_dir =r"*.jpg"
a_file = glob.glob(data_dir)[0]print(a_file, a_file.replace(".jpg",".xml"))

tree = ET.parse(a_file.replace(".jpg",".xml"))
root = tree.getroot()
size = root.find('size')
width =int(size.find('width').text)
height =int(size.find('height').text)print(f"原图宽：{width} 高：{height}")

num_classes =3
output_h = height
output_w = width
hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32)

anns =[]for obj in root.iter('object'):
    bbox = obj.find('bndbox')
    cate = obj.find('name').text
    # print(cate, bbox.find("xmin").text, bbox.find("xmax").text,#       bbox.find("ymin").text, bbox.find("ymax").text)
    xyxy =[int(bbox.find("xmin").text),int(bbox.find("ymin").text),int(bbox.find("xmax").text),int(bbox.find("ymax").text)]
    anns.append({"bbox": xyxy,'category_id':int(cate)})

num_objs =len(anns)
flipped =False#是否经过全图翻转import matplotlib.pyplot as plt
plt.figure(figsize=(19,6))
plt.ion()
plt.subplot(131)
img = plt.imread(a_file)
plt.title('Origin_img')
plt.imshow(img)for k inrange(num_objs):
    ann = anns[k]
    bbox = ann['bbox']
    cls_id = ann['category_id']if flipped:
        bbox[[0,2]]= width - bbox[[2,0]]-1# bbox[:2] = affine_transform(bbox[:2], trans_output)# 仿射变换# bbox[2:] = affine_transform(bbox[2:], trans_output)# bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, output_w - 1)#裁剪# bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, output_h - 1)
    h, w = bbox[3]- bbox[1], bbox[2]- bbox[0]if h >0and w >0:
        radius = gaussian_radius((math.ceil(h), math.ceil(w)))
        radius =max(0,int(radius))# radius = self.opt.hm_gauss if self.opt.mse_loss else radius
        ct = np.array([(bbox[0]+ bbox[2])/2,(bbox[1]+ bbox[3])/2], dtype=np.float32)
        ct_int = ct.astype(np.int32)
        plt.subplot(133)
        hm_out, gaussian = draw_umich_gaussian(hm[cls_id], ct_int, radius)
        plt.title('Umich Heatmap')# hm_out = draw_msra_gaussian(hm[cls_id], ct_int, radius)# print(hm_out.shape)# plt.title("Mara Heatmap")
        plt.text(ct[0], ct[1],f"(class:{cls_id})", c='white')
        plt.plot([bbox[0], bbox[2], bbox[2], bbox[0], bbox[0]],[bbox[1], bbox[1], bbox[3], bbox[3], bbox[1]])
        plt.imshow(hm_out)
        plt.subplot(132)
        plt.title(f'Gaussian: bbox_h={h},bbox_w={w}, radius={radius}')
        plt.imshow(gaussian)
        plt.pause(2)

参考文章

1. np.ogrid & np.mgrid 用法
2. 一维和二维高斯函数及其一阶和二阶导数