Vectorized version of argmax across 2D heatmaps

Question

Input: C=17,H,W numpy array of 17 (C=Channels) body-joint HxW heatmaps.

Desired output: Array of C=17 pairs of (Y,X) that specify the coordinates within H and W with the maximum value per each channel.

I would like to use fully vectorized solution (ideally one-liner) and replace my current per-channel solution wrapped in one ineffective "for-cycle":

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

Answer 1

Something like:

import numpy as np

C, H, W = 17, 10, 10
arr = np.random.rand(C, H, W)

out = np.array(np.unravel_index(arr.reshape(C, -1).argmax(axis=1), (H, W)))

out is shape (2, 17), which you can optionally transpose.