def gen_random_rays_from_batch_images(H, W, N_rays, images, intrinsics, c2ws, depths=None, masks=None):
    """

    :param H:
    :param W:
    :param N_rays:
    :param images: [B,3,H,W]
    :param intrinsics: [B, 3, 3]
    :param c2ws: [B, 4, 4]
    :param depths: [B,H,W]
    :param masks: [B,H,W]
    :return:
    """
    assert len(images.shape) == 4

    rays_o = []
    rays_v = []
    rays_color = []
    rays_depth = []
    rays_mask = []
    for i in range(images.shape[0]):
        sample = gen_random_rays_from_single_image(H, W, N_rays, images[i], intrinsics[i], c2ws[i],
                                                   depth=depths[i] if depths is not None else None,
                                                   mask=masks[i] if masks is not None else None)
        rays_o.append(sample['rays_o'])
        rays_v.append(sample['rays_v'])
        rays_color.append(sample['rays_color'])
        if depths is not None:
            rays_depth.append(sample['rays_depth'])
        if masks is not None:
            rays_mask.append(sample['rays_mask'])

    sample = {
        'rays_o': torch.stack(rays_o, dim=0),  # [batch, N_rays, 3]
        'rays_v': torch.stack(rays_v, dim=0),
        'rays_color': torch.stack(rays_color, dim=0),
        'rays_depth': torch.stack(rays_depth, dim=0) if depths is not None else None,
        'rays_mask': torch.stack(rays_mask, dim=0) if masks is not None else None
    }
    return sample


    :param H:
    :param W:
    :param N_rays:
    :param images: [B,3,H,W]
    :param intrinsics: [B, 3, 3]
    :param c2ws: [B, 4, 4]
    :param depths: [B,H,W]
    :param masks: [B,H,W]
    :return:
    
mv-diffusion/One-2-3-45/reconstruction/models/rays.py