From 0734dfc63c4ef9eb7a0ed4fe992ade18fedc7dad Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 7 Jul 2022 10:25:19 -0700 Subject: [PATCH] Add extra dim to bins --- pyrad/cameras/rays.py | 29 ++++++++++++++--------------- pyrad/graphs/instant_ngp.py | 2 +- pyrad/graphs/modules/ray_sampler.py | 4 ++-- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index 095ebbe89..48e5637c5 100644 --- a/pyrad/cameras/rays.py +++ b/pyrad/cameras/rays.py @@ -223,46 +223,45 @@ class RayBundle(TensorDataclass): return self.flatten()[start_idx:end_idx] def get_ray_samples( - self, bin_starts: TensorType["num_rays", "num_samples"], bin_ends: TensorType["num_rays", "num_samples"] + self, bin_starts: TensorType["num_rays", "num_samples", 1], bin_ends: TensorType["num_rays", "num_samples", 1] ) -> RaySamples: """Produces samples for each ray by projection points along the ray direction. Args: - bin_starts (TensorType["num_rays", "num_samples"]): Distance from origin to start of bin. - bin_ends (TensorType["num_rays", "num_samples"]): Distance from origin to end of bin. + bin_starts (TensorType["num_rays", "num_samples", 1]): Distance from origin to start of bin. + bin_ends (TensorType["num_rays", "num_samples", 1]): Distance from origin to end of bin. Returns: RaySamples: Samples projected along ray. """ device = self.origins.device - num_samples = bin_starts.shape[-1] valid_mask = torch.ones((bin_starts.shape), dtype=torch.bool, device=device) - dists = bin_ends - bin_starts # [N_rays, N_samples] - deltas = dists * torch.norm(self.directions[:, None, :], dim=-1) + dists = bin_ends - bin_starts # [N_rays, N_samples, 1] + deltas = dists * torch.norm(self.directions[:, :], dim=-1)[..., None, None] if is_not_none(self.camera_indices): - camera_indices = self.camera_indices.repeat(1, num_samples)[..., None] + camera_indices = self.camera_indices[..., None] else: camera_indices = None frustums = Frustums( origins=self.origins[:, None, :], # [N_rays, 1, 3] directions=self.directions[:, None, :], # [N_rays, 1, 3] - frustum_starts=bin_starts[:, :, None], # [N_rays, N_samples, 1] - frustum_ends=bin_ends[:, :, None], # [N_rays, N_samples, 1] - pixel_area=self.pixel_area[:, None, :], # [N_rays, 2, 1] + frustum_starts=bin_starts, # [N_rays, N_samples, 1] + frustum_ends=bin_ends, # [N_rays, N_samples, 1] + pixel_area=self.pixel_area[:, None, :], # [N_rays, 1, 1] ).to(device) ray_samples = RaySamples( frustums=frustums, - camera_indices=camera_indices, # [N_rays, N_samples, 1] - valid_mask=valid_mask[..., None], # [N_rays, N_samples, 1] - bin_starts=bin_starts[..., None], # [N_rays, N_samples, 1] - bin_ends=bin_ends[..., None], # [N_rays, N_samples, 1] - deltas=deltas[..., None], # [N_rays, N_samples, 1] + camera_indices=camera_indices, # [N_rays, 1, 1] + valid_mask=valid_mask, # [N_rays, N_samples, 1] + bin_starts=bin_starts, # [N_rays, N_samples, 1] + bin_ends=bin_ends, # [N_rays, N_samples, 1] + deltas=deltas, # [N_rays, N_samples, 1] ) return ray_samples diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index df8a78264..69ce8b349 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/pyrad/graphs/instant_ngp.py @@ -107,7 +107,7 @@ class NGPGraph(Graph): ts, indices = torch.sort(torch.cat([ts_uniform, ts_pdf], -2), -2) bin_starts = ts bin_ends = torch.cat([ts[..., 1:, :], ts[..., -1:, :]], dim=-2) - ray_samples = ray_bundle.get_ray_samples(bin_starts=bin_starts[..., 0], bin_ends=bin_ends[..., 0]) + ray_samples = ray_bundle.get_ray_samples(bin_starts=bin_starts, bin_ends=bin_ends) field_outputs = {} for fo_name, _ in field_outputs_pdf.items(): diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index 08cc507c5..4301fdf57 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/pyrad/graphs/modules/ray_sampler.py @@ -115,7 +115,7 @@ class SpacedSampler(Sampler): bin_lower = torch.cat([bins[..., :1], bin_centers], -1) bins = bin_lower + (bin_upper - bin_lower) * t_rand - ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1], bin_ends=bins[..., 1:]) + ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None]) return ray_samples @@ -356,6 +356,6 @@ class PDFSampler(Sampler): # Stop gradients bins = bins.detach() - ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1], bin_ends=bins[..., 1:]) + ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None]) return ray_samples