def gen_random_rays_from_single_image(H, W, N_rays, image, intrinsic, c2w, depth=None, mask=None, dilated_mask=None,
                                      importance_sample=False, h_patch_size=3):# 256 256 512 image.shape([3, 256, 256]) depth.Size([256, 256]) mask.shape.Size([256, 256]) dilated_mask=None importance_sample=None h_patch_size=3
    """
    generate random rays in world space, for a single image
    :param H:
    :param W:
    :param N_rays:
    :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

    if dilated_mask is None:# True
        dilated_mask = mask

    if not importance_sample:# False
        pixels_x = torch.randint(low=0, high=W, size=[N_rays])
        pixels_y = torch.randint(low=0, high=H, size=[N_rays])
    elif importance_sample and dilated_mask is not None:  # sample more pts in the valid mask regions
        pixels_x_1 = torch.randint(low=0, high=W, size=[N_rays // 4])
        pixels_y_1 = torch.randint(low=0, high=H, size=[N_rays // 4])

        ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                                torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
        p = torch.stack([xs, ys], dim=-1)  # H, W, 2

        try:
            p_valid = p[dilated_mask > 0]  # [num, 2]
            random_idx = torch.randint(low=0, high=p_valid.shape[0], size=[N_rays // 4 * 3])
        except:
            print("high=p_valid.shape[0] = ", p_valid.shape[0])# high=p_valid.shape[0] = 0
            print("dilated_mask.shape: ", dilated_mask.shape)#  torch.Size([256, 256])
            print("dilated_mask valid number", dilated_mask.sum())# dilated_mask valid number tensor(0.)
            
            raise ValueError("hhhh")
        p_select = p_valid[random_idx]  # [N_rays//2, 2]
        pixels_x_2 = p_select[:, 0]
        pixels_y_2 = p_select[:, 1]

        pixels_x = torch.cat([pixels_x_1, pixels_x_2], dim=0).to(torch.int64)
        pixels_y = torch.cat([pixels_y_1, pixels_y_2], dim=0).to(torch.int64)

    # - crop patch from images
    offsets = build_patch_offset(h_patch_size).to(device)
    grid_patch = torch.stack([pixels_x, pixels_y], dim=-1).view(-1, 1, 2) + offsets.float()  # [N_pts, Npx, 2]
    patch_mask = (pixels_x > h_patch_size) * (pixels_x < (W - h_patch_size)) * (pixels_y > h_patch_size) * (
            pixels_y < H - h_patch_size)  # [N_pts]
    grid_patch_u = 2 * grid_patch[:, :, 0] / (W - 1) - 1
    grid_patch_v = 2 * grid_patch[:, :, 1] / (H - 1) - 1
    grid_patch_uv = torch.stack([grid_patch_u, grid_patch_v], dim=-1)  # [N_pts, Npx, 2]
    patch_color = F.grid_sample(image[None, :, :, :], grid_patch_uv[None, :, :, :], mode='bilinear',
                                padding_mode='zeros',align_corners=True)[0]  # [3, N_pts, Npx]
    patch_color = patch_color.permute(1, 2, 0).contiguous()

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u = 2 * pixels_x / (W - 1) - 1
    ndc_v = 2 * pixels_y / (H - 1) - 1
    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)

    image = image.permute(1, 2, 0)  # H ,W, C
    color = image[(pixels_y, pixels_x)]  # N_rays, 3

    if mask is not None:
        mask = mask[(pixels_y, pixels_x)]  # N_rays
        patch_mask = patch_mask * mask  # N_rays
        mask = mask.view(-1, 1)
    else:
        mask = torch.ones([N_rays, 1])

    if depth is not None:
        depth = depth[(pixels_y, pixels_x)]  # N_rays
        depth = depth.view(-1, 1)

    intrinsic_inv = torch.inverse(intrinsic)

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

    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': rays_ndc_uv,
        'rays_color': color,
        # 'rays_depth': depth,
        'rays_mask': mask,
        'rays_norm_XYZ_cam': p,  # - XYZ_cam, before multiply depth,
        'rays_patch_color': patch_color,
        'rays_patch_mask': patch_mask.view(-1, 1)
    }

    if depth is not None:
        sample['rays_depth'] = depth

    return sample

    generate random rays in world space, for a single image
    :param H:
    :param W:
    :param N_rays:
    :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