diff --git a/.vscode/settings.json b/.vscode/settings.json index fa7771bbd..599af5b4b 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -31,6 +31,7 @@ "--max-line-length=120", "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u", "--disable=too-many-arguments", - "--disable=fixme" + "--disable=fixme", + "--disable-msg=too-many-locals" ], } \ No newline at end of file diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py index 700683fd9..4c2598704 100644 --- a/mattport/nerf/graph/nerf.py +++ b/mattport/nerf/graph/nerf.py @@ -10,7 +10,7 @@ from mattport.nerf.field_modules.ray_generator import RayGenerator from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import RGBRenderer -from mattport.nerf.sampler import UniformSampler +from mattport.nerf.sampler import PDFSampler, UniformSampler class NeRFGraph(Graph): @@ -29,7 +29,7 @@ class NeRFGraph(Graph): # samplers self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64) - self.sampler_pdf = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64) + self.sampler_pdf = PDFSampler(num_samples=64) # field self.field_coarse = NeRFField() @@ -52,7 +52,7 @@ class NeRFGraph(Graph): rgb=coarse_field_outputs.rgb, density=coarse_field_outputs.density, deltas=uniform_ray_samples.deltas ) # RendererOutputs # fine network: - pdf_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples + pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs.density) # RaySamples fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs fine_renderer_outputs = self.renderer_rgb( rgb=fine_field_outputs.rgb, density=fine_field_outputs.density, deltas=pdf_ray_samples.deltas diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 83ec23989..14ebe79fd 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,6 +2,7 @@ Collection of sampling strategies """ +from typing import Optional import torch from torch import nn from torchtyping import TensorType @@ -12,17 +13,34 @@ from mattport.structures.rays import RayBundle, RaySamples class UniformSampler(nn.Module): """Sample uniformly along a ray""" + def __init__(self, near_plane: float, far_plane: float, num_samples: int) -> None: + """ + Args: + near_plane (float): Minimum distance along ray to sample + far_plane (float): Maximum distance along ray to sample + num_samples (int): Number of samples per ray + """ + super().__init__() + + self.near_plane = near_plane + self.far_plane = far_plane + self.num_samples = num_samples + def forward( - self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int + self, + ray_bundle: RayBundle, + near_plane: Optional[float] = None, + far_plane: Optional[float] = None, + num_samples: Optional[int] = None, ) -> TensorType[..., "num_samples", 3]: """Generates position samples uniformly. Args: - camera_rays (cameras.Rays): Rays to generate samples for + ray_bundle (RayBundle): Rays to generate samples for TensorType[..., "output_dim"]: A encoded input tensor - near_plane (float): Minimum distance along ray to sample - far_plane (float): Maximum distance along ray to sample - num_samples (int): Number of samples per ray + near_plane (Optional[float]): Minimum distance along ray to sample + far_plane (Optional[float]): Maximum distance along ray to sample + num_samples (Optional[int]): Number of samples per ray Returns: RaySamples: Positions and deltas for samples along a ray @@ -31,36 +49,48 @@ class UniformSampler(nn.Module): far_plane = far_plane or self.far_plane num_samples = num_samples or self.num_samples - bins = torch.linspace(near_plane, far_plane, num_samples + 1) + num_rays = ray_bundle.origins.shape[0] - t = (bins[..., 1:] + bins[..., :-1]) / 2.0 - deltas = bins[..., 1:] - bins[..., :-1] + bins = torch.linspace(near_plane, far_plane, num_samples + 1).to(ray_bundle.origins.device) # (num_samples+1,) + bins = bins.unsqueeze(0).repeat(num_rays, 1) # (num_rays, num_samples+1) - positions = camera_rays.origin + camera_rays.direction * t + ray_samples = RaySamples( + bins=bins, + ray_bundle=ray_bundle, + ) - return RaySamples(locations=positions, deltas=deltas) + return ray_samples class PDFSampler(nn.Module): """Sample based on probability distribution""" + def __init__(self, num_samples: int) -> None: + """ + Args: + num_samples (int): Number of samples per ray + """ + super().__init__() + + self.num_samples = num_samples + def forward( self, - camera_rays: cameras.Rays, - num_samples: int, - bins: TensorType[..., "num_samples"], - weights: TensorType[..., "num_samples"], + ray_bundle: RayBundle, + coase_ray_samples: RaySamples, + density: TensorType[..., "num_samples", 1], + num_samples: Optional[int] = None, randomized: bool = True, eps: float = 1e-5, ) -> TensorType[..., "num_samples", 3]: """Generates position samples given a distribution. Args: - camera_rays (cameras.Rays): Rays to generate samples for + ray_bundle (RayBundle): Rays to generate samples for TensorType[..., "output_dim"]: A encoded input tensor - num_samples (int): Number of samples per ray bins (TensorType[..., "num_samples"]): Ray bins weights: (TensorType[..., "num_samples"]): Weights for each bin + num_samples (Optional[int]): Number of samples per ray randomized: boolean: Randomize location within each bin. Defaults to True eps: float: Small value to prevent numerical issues. Defaults to 1e-5 @@ -68,6 +98,20 @@ class PDFSampler(nn.Module): RaySamples: Positions and deltas for samples along a ray """ + num_samples = num_samples or self.num_samples + + # Calculate weight contributions along ray + # Todo(matt): This computation is duplicated + delta_density = coase_ray_samples.deltas * density[..., 0] + alphas = 1 - torch.exp(-delta_density) + + transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) + transmittance = torch.cat( + [torch.zeros((*transmittance.shape[:1], 1)).to(density.device), transmittance], axis=-1 + ) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + weights = alphas * transmittance # [..., "num_samples"] + # Add small offset to rays with zero weight to prevent NaNs weights_sum = torch.sum(weights, dim=-1, keepdim=True) padding = torch.relu(eps - weights_sum) @@ -75,38 +119,33 @@ class PDFSampler(nn.Module): weights_sum += padding pdf = weights / weights_sum - cdf = torch.min(1, torch.cumsum(pdf, dim=-1)) + cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1)) cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1) if randomized: - u = torch.rand(size=(*cdf.shape[:-1], num_samples)) + u = torch.rand(size=(*cdf.shape[:-1], num_samples), device=cdf.device) else: - u = torch.linspace(0.0, 1.0, steps=num_samples) + u = torch.linspace(0.0, 1.0, steps=num_samples, device=cdf.device) u = torch.expand(size=(*cdf.shape[:-1], num_samples)) - # u = u.contigous() + # u = u.contiguous() indicies = torch.searchsorted(cdf, u, right=True) below = torch.max(torch.zeros_like(indicies), indicies - 1) - above = torch.min(cdf.shape[-1] - 1, indicies) + above = torch.min(cdf.shape[-1] - torch.ones_like(indicies), indicies) indicies_g = torch.stack([below, above], -1) - cdf_g = torch.gather(cdf, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2) - bins_g = torch.gather(bins, indicies_g, axis=-1, batch_dims=len(indicies_g.shape) - 2) + matched_shape = (indicies_g.shape[0], indicies_g.shape[1], cdf.shape[-1]) + cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) + bins_g = torch.gather(coase_ray_samples.bins.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) denom = cdf_g[..., 1] - cdf_g[..., 0] denom = torch.where(denom < eps, torch.ones_like(denom), denom) t = (u - cdf_g[..., 0]) / denom samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) - deltas = samples[..., 1:] - samples[..., :-1] - - positions = camera_rays.origin + camera_rays.direction * t - - num_rays = ray_bundle.origins.shape[0] ray_samples = RaySamples( - t_min=torch.ones_like(ray_bundle.origins[:, 0]) * near_plane, - t_max=torch.ones_like(ray_bundle.origins[:, 0]) * far_plane, - ts=bins.unsqueeze(0).repeat(num_rays, 1), # (num_rays, num_samples) + bins=samples, ray_bundle=ray_bundle, ) + return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 4b72a860d..fed7d3f4d 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -72,14 +72,10 @@ class RaySamples: def __init__( self, - t_min: TensorType["num_rays"], - t_max: TensorType["num_rays"], - ts: TensorType["num_rays", "num_samples+1"], + bins: TensorType["num_rays", "num_samples+1"], ray_bundle: RayBundle, ) -> None: - self.t_min = t_min - self.t_max = t_max - self.ts = ts + self.bins = bins self.positions = self.get_positions(ray_bundle) self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1) self.deltas = self.get_deltas() @@ -87,11 +83,11 @@ class RaySamples: def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]: """Returns positions.""" # TODO(ethan): check with Matt about this - t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2 # (num_rays, num_samples) + t_mid = (self.bins[:, 1:] + self.bins[:, :-1]) / 2 # (num_rays, num_samples) return ray_bundle.origins[:, None] + t_mid[:, :, None] * ray_bundle.directions[:, None] def get_deltas(self) -> TensorType["num_samples"]: """Returns deltas.""" - deltas = self.ts[..., 1:] - self.ts[..., :-1] + deltas = self.bins[..., 1:] - self.bins[..., :-1] # TODO(ethan): check this with Matt return deltas