import torch


def normalize_coords(coords_2D, batch_size, height, width):
    """
        Normalize 2D keypoint coordinates to lie in the range [-1, 1].

        Args:
            coords_2D (torch.tensor): 2D image coordinates store in order (u, v), (Bx2xN).
            batch_size (int): batch size.
            height (int): image height.
            width (int): image width.

        Returns:
            coords_2D_norm (torch.tensor): coordinates normalized to range [-1, 1] and stored in order (u, v), (BxNx2).
     """
    u_norm = (2 * coords_2D[:, 0, :].reshape(batch_size, -1) / (width - 1)) - 1
    v_norm = (2 * coords_2D[:, 1, :].reshape(batch_size, -1) / (height - 1)) - 1

    return torch.stack([u_norm, v_norm], dim=2)