寻求17通道人体关节热力图的全矢量化argmax实现方案

输入为形状(17, H, W)的numpy数组，包含17个通道的人体关节H×W热力图；需要输出包含17组(Y,X)坐标的数组，每组对应单个通道中数值最大的位置。要求用全矢量化方案（理想为单行代码）替代低效的逐通道循环实现：

kpts = [np.array(unravel_index(j.argmax(), j.shape)) for j in input]

一行代码完成全矢量化处理：

kpts = np.array(np.unravel_index(input.argmax(axis=(1,2)), input.shape[1:])).T

• input.argmax(axis=(1,2))：一次性对所有通道的空间维度（H和W）取最大值的扁平索引，得到长度为17的一维数组，每个元素对应一个通道的最大元素在H×W平面上的位置索引。
• np.unravel_index(..., input.shape[1:])：将所有扁平索引批量转换为对应的(Y,X)坐标，返回两个分别存储Y、X坐标的一维数组。
• .T：转置后得到(17, 2)形状的数组，结构与原循环输出完全一致，每行对应一组(Y,X)坐标。

内容的提问来源于stack exchange，提问作者Petr Elias