import torch


def point_cloud_from_depth(depth, K):
    """Transform depth image pixels selected in mask into point cloud in camera frame
    depth: [N, H, W] depth image
    mask: [N, H, W] mask of pixels to transform into point cloud
    K: [N, 3, 3] Intrinsic camera matrix
    returns: [N, 3, K]
    """
    batch_size, H, W = depth.shape
    # Create 3D grid data
    u_img_range = torch.arange(0, W, device=depth.device)
    v_img_range = torch.arange(0, H, device=depth.device)
    u_grid, v_grid = torch.meshgrid(u_img_range, v_img_range)
    u_grid = u_grid.t()[None, :].repeat(batch_size, 1, 1)
    v_grid = v_grid.t()[None, :].repeat(batch_size, 1, 1)
    u_img, v_img, d = (
        u_grid.reshape(batch_size, -1),
        v_grid.reshape(batch_size, -1),
        depth.reshape(batch_size, -1),
    )

    # homogenuous coordinates
    uv = torch.stack((u_img, v_img, torch.ones_like(u_img)), dim=1).float()

    # get the unscaled position for each of the points in the image frame
    unscaled_points = torch.linalg.inv(K) @ uv

    # scale points by their depth value
    return d[:, None] * unscaled_points