| 22 | |
| 23 | # Hierarchical sampling (section 5.2) |
| 24 | def sample_pdf(bins, weights, N_samples, det=False, pytest=False): |
| 25 | # Get pdf |
| 26 | weights = weights + 1e-5 # prevent nans |
| 27 | pdf = weights / torch.sum(weights, -1, keepdim=True) |
| 28 | cdf = torch.cumsum(pdf, -1) |
| 29 | cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # (batch, len(bins)) |
| 30 | |
| 31 | # Take uniform samples |
| 32 | if det: |
| 33 | u = torch.linspace(0., 1., steps=N_samples) |
| 34 | u = u.expand(list(cdf.shape[:-1]) + [N_samples]) |
| 35 | else: |
| 36 | u = torch.rand(list(cdf.shape[:-1]) + [N_samples]) |
| 37 | |
| 38 | # Pytest, overwrite u with numpy's fixed random numbers |
| 39 | if pytest: |
| 40 | np.random.seed(0) |
| 41 | new_shape = list(cdf.shape[:-1]) + [N_samples] |
| 42 | if det: |
| 43 | u = np.linspace(0., 1., N_samples) |
| 44 | u = np.broadcast_to(u, new_shape) |
| 45 | else: |
| 46 | u = np.random.rand(*new_shape) |
| 47 | u = torch.Tensor(u) |
| 48 | |
| 49 | # Invert CDF |
| 50 | u = u.contiguous() |
| 51 | inds = torch.searchsorted(cdf.detach(), u, right=True) |
| 52 | below = torch.max(torch.zeros_like(inds-1), inds-1) |
| 53 | above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds) |
| 54 | inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2) |
| 55 | |
| 56 | matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]] |
| 57 | cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g) |
| 58 | bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g) |
| 59 | |
| 60 | denom = (cdf_g[...,1]-cdf_g[...,0]) |
| 61 | denom = torch.where(denom<1e-5, torch.ones_like(denom), denom) |
| 62 | t = (u-cdf_g[...,0])/denom |
| 63 | samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0]) |
| 64 | |
| 65 | return samples |
| 66 | |
| 67 | def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=False): |
| 68 | ''' |