import torch


def predict_keypoints(rois, heatmap):
    """
    Arguments:
        rois (FloatTensor, fp32): [R, >=5] (batch_ix, x1, y1, x2, y2, ...)
        heatmap (FloatTensor, fp32): [R, K, h, w]

    Returns:
        keyps (FloatTensor): [R, K, 3] (x, y, score)
    """
    rois_tl_x = rois[:, 1]
    rois_tl_y = rois[:, 2]
    rois_w = torch.clamp(rois[:, 3] - rois[:, 1] + 1, min=1)
    rois_h = torch.clamp(rois[:, 4] - rois[:, 2] + 1, min=1)

    R, K, H, W = heatmap.shape
    heatmap = heatmap.reshape(R, K, -1)
    score, index = heatmap.max(dim=2)
    x = (index % W).float()
    y = (index // W).float()
    x = (x + 0.5) / W * rois_w[:, None] + rois_tl_x[:, None]
    y = (y + 0.5) / H * rois_h[:, None] + rois_tl_y[:, None]
    return torch.stack([x, y, score], dim=2)