def gen_random_rays_of_patch_from_single_image(H, W, N_rays, num_neighboring_pts, patch_size,
                                               image, intrinsic, c2w, depth=None, mask=None):
    """
    generate random rays in world space, for a single image
    sample rays from local patches
    :param H:
    :param W:
    :param N_rays: the number of center rays of patches
    :param image: [3, H, W]
    :param intrinsic: [3,3]
    :param c2w: [4,4]
    :param depth: [H, W]
    :param mask: [H, W]
    :return:
    """
    device = image.device
    patch_radius_max = patch_size // 2

    unit_u = 2 / (W - 1)
    unit_v = 2 / (H - 1)

    pixels_x_center = torch.randint(low=patch_size, high=W - patch_size, size=[N_rays])
    pixels_y_center = torch.randint(low=patch_size, high=H - patch_size, size=[N_rays])

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u_center = 2 * pixels_x_center / (W - 1) - 1
    ndc_v_center = 2 * pixels_y_center / (H - 1) - 1
    ndc_uv_center = torch.stack([ndc_u_center, ndc_v_center], dim=-1).view(-1, 2).float().to(device)[:, None,
                    :]  # [N_rays, 1, 2]

    shift_u, shift_v = torch.rand([N_rays, num_neighboring_pts, 1]), torch.rand(
        [N_rays, num_neighboring_pts, 1])  # uniform distribution of [0,1)
    shift_u = 2 * (shift_u - 0.5)  # mapping to [-1, 1)
    shift_v = 2 * (shift_v - 0.5)

    # - avoid sample points which are too close to center point
    shift_uv = torch.cat([(shift_u * patch_radius_max) * unit_u, (shift_v * patch_radius_max) * unit_v],
                         dim=-1)  # [N_rays, num_npts, 2]
    neighboring_pts_uv = ndc_uv_center + shift_uv  # [N_rays, num_npts, 2]

    sampled_pts_uv = torch.cat([ndc_uv_center, neighboring_pts_uv], dim=1)  # concat the center point

    # sample the gts
    color = F.grid_sample(image[None, :, :, :], sampled_pts_uv[None, :, :, :], mode='bilinear',
                          align_corners=True)[0]  # [3, N_rays, num_npts]
    depth = F.grid_sample(depth[None, None, :, :], sampled_pts_uv[None, :, :, :], mode='bilinear',
                          align_corners=True)[0]  # [1, N_rays, num_npts]

    mask = F.grid_sample(mask[None, None, :, :].to(torch.float32), sampled_pts_uv[None, :, :, :], mode='nearest',
                         align_corners=True).to(torch.int64)[0]  # [1, N_rays, num_npts]

    intrinsic_inv = torch.inverse(intrinsic)

    sampled_pts_uv = sampled_pts_uv.view(N_rays * (1 + num_neighboring_pts), 2)
    color = color.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 3)
    depth = depth.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 1)
    mask = mask.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 1)

    pixels_x = (sampled_pts_uv[:, 0] + 1) * (W - 1) / 2
    pixels_y = (sampled_pts_uv[:, 1] + 1) * (H - 1) / 2
    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float().to(device)  # N_rays*num_pts, 3
    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays*num_pts, 3
    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays*num_pts, 3
    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays*num_pts, 3
    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays*num_pts, 3

    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': sampled_pts_uv,
        'rays_color': color,
        'rays_depth': depth,
        'rays_mask': mask,
        # 'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
    }

    return sample

    generate random rays in world space, for a single image
    sample rays from local patches
    :param H:
    :param W:
    :param N_rays: the number of center rays of patches
    :param image: [3, H, W]
    :param intrinsic: [3,3]
    :param c2w: [4,4]
    :param depth: [H, W]
    :param mask: [H, W]
    :return:
    
mv-diffusion/One-2-3-45-by-view/reconstruction/models/rays.py