import torch


def heatmap_to_ann(heatmap_batch):
  """ Convert heatmap into series of X,Y coordinate by applying argmax.
  Args:
    heatmap_batch (torch.tensor)

  Returns: Integer coordinates (x, y)
  """
  n_batch, n_lm, img_w, img_h = heatmap_batch.shape
  x = torch.zeros([n_batch, n_lm])
  y = torch.zeros([n_batch, n_lm])

  for n in range(n_batch):
    for i in range(n_lm):
      raw_idx = heatmap_batch[n, i, :, :].argmax()
      y[n, i] = raw_idx // img_h
      x[n, i] = raw_idx - (y[n, i] * img_h)
  
  return x.int(), y.int()