tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Add extra dim to bins

Browse Source
This commit is contained in:
Matt Tancik 2022-07-07 10:25:19 -07:00 committed by Matthew Tancik
parent 44c5e3264c
commit 0734dfc63c
3 changed files with 17 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
pyrad/cameras/rays.py
View File

@ -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

2
pyrad/graphs/instant_ngp.py
View File

@ -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():

4
pyrad/graphs/modules/ray_sampler.py
View File

@ -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