diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 931e2bec2..5eb4f7fe7 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -179,7 +179,7 @@ class NeRFGraph(Graph): depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts) # pdf sampling - ray_samples_pdf = self.sampler_pdf(ray_samples_uniform, weights_coarse) + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 9269322d9..33a85921d 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -94,10 +94,7 @@ class UniformSampler(Sampler): else: ts = (bins[:, 1:] + bins[:, :-1]) / 2 - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples @@ -129,6 +126,7 @@ class PDFSampler(Sampler): @torch.no_grad() def generate_ray_samples( self, + ray_bundle: RayBundle = None, ray_samples: RaySamples = None, weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, @@ -189,9 +187,6 @@ class PDFSampler(Sampler): # Stop gradients ts = ts.detach() - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_samples.ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 56a8134f7..7593041ea 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -21,6 +21,50 @@ class PointSamples: valid_mask: TensorType[..., 1] = None +@dataclass +class RaySamples: + """Samples along a ray""" + + positions: TensorType[..., 3] = None + directions: TensorType[..., 3] = None + camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[..., 1] = None + ts: TensorType[..., 1] = None + deltas: TensorType[..., 1] = None + + def to_point_samples(self) -> PointSamples: + """Convert to PointSamples instance and return.""" + # TODO: make this more interpretable + return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) + + def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: + """Return weights based on predicted densities + + Args: + densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray + + Returns: + TensorType[..., "num_samples"]: Weights for each sample + """ + + delta_density = self.deltas * densities[..., 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(densities.device), transmittance], axis=-1 + ) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + + weights = alphas * transmittance # [..., "num_samples"] + + return weights + + def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: + """Sets valid mask""" + self.valid_mask = valid_mask + + @dataclass class RayBundle: """_summary_ @@ -86,6 +130,38 @@ class RayBundle: valid_mask=self.valid_mask[valid_mask] if is_not_none(self.valid_mask) else None, ) + def get_ray_samples(self, ts: TensorType["num_rays", "num_samples"]) -> RaySamples: + """ + Args: + ts (TensorType["num_rays", "num_samples"]): _description_ + + Returns: + RaySamples: _description_ + """ + positions = self.origins[:, None] + ts[:, :, None] * self.directions[:, None] + directions = self.directions.unsqueeze(1).repeat(1, positions.shape[1], 1) + valid_mask = torch.ones_like(ts, dtype=torch.bool) + + dists = ts[..., 1:] - ts[..., :-1] + dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] + deltas = dists * torch.norm(self.directions[..., None, :], dim=-1) + + if self.camera_indices is not None: + camera_indices = self.camera_indices.unsqueeze(1).repeat(1, positions.shape[1]) + else: + camera_indices = None + + ray_samples = RaySamples( + positions=positions, + directions=directions, + camera_indices=camera_indices, + valid_mask=valid_mask, + ts=ts, + deltas=deltas, + ) + + return ray_samples + @dataclass class CameraRayBundle: @@ -131,69 +207,3 @@ class CameraRayBundle: directions=self.directions.view(-1, 3)[start_idx:end_idx], camera_indices=camera_indices, ) - - -class RaySamples: - """_summary_""" - - def __init__( - self, - ts: TensorType["num_rays", "num_samples"], - ray_bundle: RayBundle, - ) -> None: - self.ray_bundle = ray_bundle - self.ts = ts - 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() - self.valid_mask = torch.ones_like(self.ts, dtype=torch.bool) - - def to_point_samples(self) -> PointSamples: - """Convert to PointSamples instance and return.""" - # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) - - def get_camera_indices(self): - """Returns camera indices.""" - assert not isinstance(self.ray_bundle.camera_indices, type(None)), "Camera indices cannot be None here." - camera_indices = self.ray_bundle.camera_indices.unsqueeze(1).repeat(1, self.positions.shape[1]) - return camera_indices - - def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]: - """Returns positions.""" - return ray_bundle.origins[:, None] + self.ts[:, :, None] * ray_bundle.directions[:, None] - - def get_deltas(self) -> TensorType[..., "num_samples"]: - """Returns deltas.""" - dists = self.ts[..., 1:] - self.ts[..., :-1] - dist_inf = 1e10 * torch.ones_like(dists[..., -1:]) - dists = torch.cat([dists, dist_inf], -1) # [N_rays, N_samples] - deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1) - return deltas - - def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: - """Return weights based on predicted densities - - Args: - densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray - - Returns: - TensorType[..., "num_samples"]: Weights for each sample - """ - - delta_density = self.deltas * densities[..., 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(densities.device), transmittance], axis=-1 - ) - transmittance = torch.exp(-transmittance) # [..., "num_samples"] - - weights = alphas * transmittance # [..., "num_samples"] - - return weights - - def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: - """Sets valid mask""" - self.valid_mask = valid_mask diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index 92a6848bb..df88b889c 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -47,7 +47,7 @@ def test_pdf_sampler(): # Just check that it doesn't crash. Need to add some actual tests. pdf_sampler = PDFSampler(num_samples) - pdf_sampler(coarse_ray_samples, weights, num_samples) + pdf_sampler(ray_bundle, coarse_ray_samples, weights, num_samples) # TODO Tancik: Add more precise tests