import torch


def normalize_coords(coords_2D, width, height):
    """Normalizes coords_2D (BW x N x 2) in pixel coordinates to be within [-1, 1]
    Args:
        coords_2D (torch.tensor): (b*w,N,2)
        width (float): width of the image in pixels
        height (float): height of the image in pixels
    Returns:
        torch.tensor: (b*w,N,2) coordinates normalized to be within [-1,1]
    """
    batch_size = coords_2D.size(0)
    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)