热图回归是关键点检测中常见的解决方案，一般采用将关键点通过高斯分布的方式扩散到周围的其他像素/体素的方式构建热图，并通过让模型直接预测热图->通过热图极值点和局部拟合的方式还原关键点来达到预测关键点的目的。热图回归部分已经有不少的资料进行解读，然而解码可能是因为太简单了，网上没太找到相关的资料，于是写了这篇博客大概记录一下。
热图回归的公式一般表述如下，对于图像上的任意一个点 $(x, y, z)$，其在对应第 $k$ 个关键点的热图 $H_{k}$ 上的值可以由以下公式计算，其中 $\sigma$ 是标准差，决定了热图高亮区域的大小，是经验值，建议在实际使用中可视化热图来确定：
$$ H_{k}(x,y,z)=exp(-\frac{(x-x{k})^2+(y-y_{k})^2+(z-z_{k})^2}{2\sigma^2}) $$
也可以使用各向异性的高斯分布来表示图像中各轴spacing的差异或者构造更符合实际的热力图，公式如下：
$$ V(x,y,z)=exp(-(\frac{(x-x_{k})^2}{2\sigma_{x}^2}+\frac{(y-y_{k}^2)}{2\sigma_{y}^2}+\frac{(z-z_{k})^2}{2\sigma_{z}^2 })) $$
解码一般分为两步：1.找极大值 2.精度优化。
找极大值有两个情况，一种是确定单峰值，如人体检测中确定人只有一个脑/左肺/右肺，这个时候直接argmax是最高效的做法：
import numpy as np
def get_max_preds_3d(batch_heatmaps_3d):
assert batch_heatmaps_3d.ndim == 5,
batch_size, num_keypoints, depth, height, width = batch_heatmaps_3d.shape
heatmaps_reshaped = batch_heatmaps_3d.reshape((batch_size, num_keypoints, -1))
idx = np.argmax(heatmaps_reshaped, 2)
maxvals = np.amax(heatmaps_reshaped, 2)
maxvals = maxvals.reshape((batch_size, num_keypoints, 1))
preds = np.zeros((batch_size, num_keypoints, 3), dtype=np.float32)
preds[:, :, 0] = idx % width
preds[:, :, 1] = (idx // width) % height
preds[:, :, 2] = idx // (width * height)
return preds, maxvals
第二种是存在多峰值，这个时候需要进行峰值查找，一般做法为NMS+局部拟合
NMS的作用是确保一个峰值只被检测一次，一个体素只有当其值大于其卷积核内所有其他体素时，才被视为峰顶，以此获得粗略的坐标估计。局部拟合是假设每个峰值周围大致符合高斯分布，通过计算质心的方式对点进行微调，做这步的原因是我们的热图大部分情况下都不会在原尺寸下进行预测，而是先进行下采样，这一步中存在精度损失需要补偿。
def nms_3d(heatmap, kernel_size=3, threshold=0.5):
"""
3D非极大值抑制
Args:
heatmap: 热图 [B, 1, D, H, W]
kernel_size: NMS核大小
threshold: 阈值
Returns:
peaks: 关键点坐标 [B, N, 3]，坐标范围[0,1]
"""
# 确保输入是float类型
heatmap = heatmap.float()
# 创建padding
pad = (kernel_size - 1) // 2
padded_heatmap = F.pad(heatmap, (pad, pad, pad, pad, pad, pad), mode='reflect')
# 创建最大池化层（最大值滤波）
max_pool = nn.MaxPool3d(kernel_size=kernel_size, stride=1, padding=0)
# 获取局部最大值
local_max = max_pool(padded_heatmap)
# 创建掩码
mask = (heatmap == local_max) & (heatmap > threshold)
# 获取每个批次的关键点
batch_peaks = []
batch_scores = []
for b in range(heatmap.shape[0]):
# 获取当前批次的热图
curr_heatmap = heatmap[b, 0]
curr_mask = mask[b, 0]
# 获取关键点坐标
peak_coords = torch.nonzero(curr_mask).float()
# 如果没有检测到关键点，返回空列表
if len(peak_coords) == 0:
batch_peaks.append(torch.zeros((0, 3), device=heatmap.device))
batch_scores.append(torch.zeros((0), device=heatmap.device))
continue
# 获取关键点分数
peak_scores = curr_heatmap[curr_mask]
# 按分数排序
sorted_scores, sort_idx = torch.sort(peak_scores, descending=True)
peak_coords = peak_coords[sort_idx]
# 归一化坐标到[0,1]范围
peak_coords[:, 0] = peak_coords[:, 0] / (heatmap.shape[2] - 1)
peak_coords[:, 1] = peak_coords[:, 1] / (heatmap.shape[3] - 1)
peak_coords[:, 2] = peak_coords[:, 2] / (heatmap.shape[4] - 1)
batch_peaks.append(peak_coords)
batch_scores.append(sorted_scores)
# 填充到相同长度
max_peaks = max(len(peaks) for peaks in batch_peaks)
padded_peaks = []
padded_scores = []
for peaks, scores in zip(batch_peaks, batch_scores):
if len(peaks) < max_peaks:
padding = torch.zeros((max_peaks - len(peaks), 3), device=heatmap.device)
peaks = torch.cat([peaks, padding], dim=0)
score_padding = torch.zeros((max_peaks - len(scores)), device=heatmap.device)
scores = torch.cat([scores, score_padding], dim=0)
padded_peaks.append(peaks)
padded_scores.append(scores)
else:
padded_peaks.append(peaks)
padded_scores.append(scores)
return torch.stack(padded_peaks), torch.stack(padded_scores)
局部拟合：
def gaussian_fit(heatmap, peaks, window_size=3):
"""
使用高斯拟合优化关键点位置
Args:
heatmap: 热图 [B, 1, D, H, W]
peaks: 初始关键点坐标 [B, N, 3]
window_size: 拟合窗口大小
Returns:
refined_peaks: 优化后的关键点坐标 [B, N, 3]
"""
B, N = peaks.shape[:2]
refined_peaks = []
for b in range(B):
batch_peaks = []
for n in range(N):
# 获取当前关键点坐标
x, y, z = peaks[b, n]
# 如果坐标为0（填充值），跳过
if torch.all(peaks[b, n] == 0):
batch_peaks.append(peaks[b, n])
continue
# 转换回像素坐标
x = int(x * (heatmap.shape[4] - 1))
y = int(y * (heatmap.shape[3] - 1))
z = int(z * (heatmap.shape[2] - 1))
# 提取局部窗口
pad = window_size // 2
local_heatmap = heatmap[b, 0,
max(0, z-pad):min(heatmap.shape[2], z+pad+1),
max(0, y-pad):min(heatmap.shape[3], y+pad+1),
max(0, x-pad):min(heatmap.shape[4], x+pad+1)]
# 计算质心
z_coords = torch.arange(local_heatmap.shape[0], device=heatmap.device)
y_coords = torch.arange(local_heatmap.shape[1], device=heatmap.device)
x_coords = torch.arange(local_heatmap.shape[2], device=heatmap.device)
z_grid, y_grid, x_grid = torch.meshgrid(z_coords, y_coords, x_coords, indexing='ij')
total_weight = local_heatmap.sum()
if total_weight > 0:
z_refined = (z_grid * local_heatmap).sum() / total_weight
y_refined = (y_grid * local_heatmap).sum() / total_weight
x_refined = (x_grid * local_heatmap).sum() / total_weight
# 转换回相对坐标
z_refined = (z_refined + max(0, z-pad)) / (heatmap.shape[2] - 1)
y_refined = (y_refined + max(0, y-pad)) / (heatmap.shape[3] - 1)
x_refined = (x_refined + max(0, x-pad)) / (heatmap.shape[4] - 1)
batch_peaks.append(torch.tensor([z_refined, y_refined, x_refined], device=heatmap.device))
else:
batch_peaks.append(peaks[b, n])
# 如果 batch_peaks 为空，添加一个零张量
if len(batch_peaks) == 0:
batch_peaks.append(torch.zeros(3, device=heatmap.device))
refined_peaks.append(torch.stack(batch_peaks))
return torch.stack(refined_peaks)