import torch


def project_points(points_3d, K, pose, z_min=None):
    """
    Args:
        points_3d: BxPx3
        K: Bx3x3
        pose: Bx3x4
        z_min: prevent zero devision, eg. 0.1
    Returns:
        projected 2d points: BxPx2
    """
    assert K.shape[-2:] == (3, 3)
    assert pose.shape[-2:] == (3, 4)
    batch_size = points_3d.shape[0]
    n_points = points_3d.shape[1]
    device = points_3d.device
    if points_3d.shape[-1] == 3:
        points_3d = torch.cat((points_3d, torch.ones(batch_size, n_points, 1).to(device)), dim=-1)
    P = K @ pose[:, :3]
    suv = (P.unsqueeze(1) @ points_3d.unsqueeze(-1)).squeeze(-1)  # Bx1x3x4 @ BxPx4x1 -> BxPx3
    if z_min is not None:
        z = suv[..., -1]
        suv[..., -1] = torch.max(torch.ones_like(z) * z_min, z)  # eg. z_min=0.1
    suv = suv / suv[..., [-1]]
    return suv[..., :2]