pyskl/datasets/pipelines/heatmap_related.py

class GeneratePoseTarget:

首先看def __call__(self, results):

def __call__(self, results):
        heatmap = self.gen_an_aug(results)
        key = 'heatmap_imgs' if 'imgs' in results else 'imgs'

        if self.double:
            indices = np.arange(heatmap.shape[1], dtype=np.int64)
            left, right = (self.left_kp, self.right_kp) if self.with_kp else (self.left_limb, self.right_limb)
            for l, r in zip(left, right):  # noqa: E741
                indices[l] = r
                indices[r] = l
            heatmap_flip = heatmap[..., ::-1][:, indices]
            heatmap = np.concatenate([heatmap, heatmap_flip])
        results[key] = heatmap
        return results

第一句

heatmap = self.gen_an_aug(results)

这里是关键

进入gen_an_aug函数经过分析

调用 self.generate_heatmap 方法表示生成当前帧的热图，并存储在 ret 中，意味着每一帧都会生成一个对应的热图。

经过分析发现gen_an_aug函数的返回值的shape是(480, 17, 64, 64)

480：表示有 480 帧

17：表示每帧中有 17 个关键点

64：表示热图的高度

64：表示热图的宽度

第二句

key = 'heatmap_imgs' if 'imgs' in results else 'imgs'

• 如果 results 字典中存在 imgs 键，则使用 heatmap_imgs 作为新键；否则，使用 imgs 作为键。
• 经过分析发现results 字典中不存在imgs 键，于是key = 'imgs'

言外之意

key = 'imgs'

第三句

由于self.double默认是False

double (bool): Output both original heatmaps and flipped heatmaps.
            Default: False.

于是只需要看

        results[key] = heatmap
        return results

将生成的热图存储到 results 字典中，使用之前确定的键名。

https://github.com/kennymckormick/pyskl/blob/main/pyskl/datasets/pipelines/heatmap_related.py