How to get the coordinates from the heatmaps？

[SunnyDp]

Thanks to your great job! In the function “find_tensor_peak_batch”, the keypoints coordinats can be obtained by using "argmax", but in your implementation, the function "affine_grid" which always shows up in Spatial Transform Network is used to get the coordinates. I wonder why this step is necessary, looking forward your reply.

[D-X-Y]

Thanks for your interest in our work! "argmax" is non-differentiable and does not allow gradients to go through the coordinates. We use affine_grid to replace this hard operation (argmax) by a soft manner, which allows gradients through the coordinates.