from typing import Union
import torch


def cam_init2orig(cam, scale: Union[float, torch.Tensor], start_pt: torch.Tensor, N=224):
    """
    Args:
        cam (bs, 3): (s, tx, ty)
        scale (bs,): scale = resize_h / orig_h
        start_pt (bs, 2): (lt_x, lt_y)
        N (int): hmr_image_size (224) or IMG_SIZE

    Returns:
        cam_orig (bs, 3): (s, tx, ty), camera in original image coordinates.

    """

    # This is camera in crop image coord.
    cam_crop = torch.cat(
        [N * cam[:, 0:1] * 0.5, cam[:, 1:] + (2. / cam[:, 0:1]) * 0.5],
        dim=1
    )

    # This is camera in orig image coord
    cam_orig = torch.cat(
        [cam_crop[:, 0:1] / scale, cam_crop[:, 1:] + (start_pt - N) / cam_crop[:, 0:1]],
        dim=1
    )

    return cam_orig