表示生成当前帧的热图,并存储在 ret 中,意味着每一帧都会生成一个对应的热图。
这里的i表示的是第i帧
进入generate_heatmap函数源码
发现注释确实写了这是处理的单独一帧,生成一帧伪热图
由于self.with_kp是True,而self.with_limb是False
于是只需要看下面
def generate_heatmap(self, arr, kps, max_values):
"""Generate pseudo heatmap for all keypoints and limbs in one frame (if
needed).
Args:
arr (np.ndarray): The array to store the generated heatmaps. Shape: V * img_h * img_w.
kps (np.ndarray): The coordinates of keypoints in this frame. Shape: M * V * 2.
max_values (np.ndarray): The confidence score of each keypoint. Shape: M * V.
Returns:
np.ndarray: The generated pseudo heatmap.
"""
if self.with_kp:
num_kp = kps.shape[1]
for i in range(num_kp):
self.generate_a_heatmap(arr[i], kps[:, i], max_values[:, i])
self.generate_a_heatmap(arr[i], kps[:, i], max_values[:, i])
表示某一帧第i个关键点生成对应的关键点的热图
进入generate_a_heatmap函数源码
发现注释确实写了这是处理的单独一帧,单独一个关键点生成伪热图
def generate_a_heatmap(self, arr, centers, max_values):
"""Generate pseudo heatmap for one keypoint in one frame.
Args:
arr (np.ndarray): The array to store the generated heatmaps. Shape: img_h * img_w.
centers (np.ndarray): The coordinates of corresponding keypoints (of multiple persons). Shape: M * 2.
max_values (np.ndarray): The max values of each keypoint. Shape: M.
Returns:
np.ndarray: The generated pseudo heatmap.
"""
sigma = self.sigma
img_h, img_w = arr.shape
for center, max_value in zip(centers, max_values):
if max_value < EPS:
continue
mu_x, mu_y = center[0], center[1]
st_x = max(int(mu_x - 3 * sigma), 0)
ed_x = min(int(mu_x + 3 * sigma) + 1, img_w)
st_y = max(int(mu_y - 3 * sigma), 0)
ed_y = min(int(mu_y + 3 * sigma) + 1, img_h)
x = np.arange(st_x, ed_x, 1, np.float32)
y = np.arange(st_y, ed_y, 1, np.float32)
# if the keypoint not in the heatmap coordinate system
if not (len(x) and len(y)):
continue
y = y[:, None]
patch = np.exp(-((x - mu_x)**2 + (y - mu_y)**2) / 2 / sigma**2)
patch = patch * max_value
arr[st_y:ed_y, st_x:ed_x] = np.maximum(arr[st_y:ed_y, st_x:ed_x], patch)
pyskl/pyskl/datasets/pipelines/heatmap_related.py at main · kennymckormick/pyskl · GitHub