tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Disable occlusion sampling and speed up eval (#488)

Browse Source 
disable occlusion sampling and speed up eval
This commit is contained in:
Matthew Tancik 2022-09-19 17:38:47 -07:00 committed by GitHub
parent 4250161922
commit 96e1f2a698
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 11 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nerfactory/cameras/cameras.py
View File

@ -232,17 +232,17 @@ class Cameras:
dx = torch.sqrt(torch.sum((directions - directions_stack[1]) ** 2, dim=-1))
dy = torch.sqrt(torch.sum((directions - directions_stack[2]) ** 2, dim=-1))
pixel_area = dx * dy
pixel_area = (dx * dy)[..., None]
if not isinstance(camera_indices, torch.Tensor):
ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to((self._num_cameras)).to(self.device)
ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to(pixel_area.shape).to(self.device)
else:
ray_bundle_camera_indices = camera_indices
return RayBundle(
origins=origins,
directions=directions,
pixel_area=pixel_area[..., None],
pixel_area=pixel_area,
camera_indices=ray_bundle_camera_indices,
)

2
nerfactory/datamanagers/base.py
View File

@ -332,6 +332,6 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
for camera_ray_bundle, batch in self.eval_dataloader:
assert camera_ray_bundle.camera_indices is not None
image_idx = int(camera_ray_bundle.camera_indices[0, 0])
image_idx = int(camera_ray_bundle.camera_indices[0, 0, 0])
return image_idx, camera_ray_bundle, batch
raise ValueError("No more eval images")

1
nerfactory/datamanagers/dataloaders.py
View File

@ -156,7 +156,6 @@ class EvalDataloader(DataLoader):
image_idx: Camera image index
"""
ray_bundle = self.cameras.generate_rays(camera_indices=image_idx)
ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int()
batch = self.input_dataset[image_idx]
batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
return ray_bundle, batch

15
nerfactory/models/modules/ray_sampler.py
View File

@ -347,6 +347,12 @@ class VolumetricSampler(Sampler):
"""Sampler inspired by the one proposed in the Instant-NGP paper.
This sampler does ray-box AABB test, ray samples generation and density check, all together.
Args:
aabb: Bounding box of the scene.
density_fn: Function that takes a tensor of points and returns a tensor of densities.
grid_resolution: Resolution of the density grid.
num_samples: Number of samples per ray.
"""
def __init__(
@ -356,14 +362,7 @@ class VolumetricSampler(Sampler):
grid_resolution: int = 128,
num_samples: int = 1024,
) -> None:
"""Init.
Args:
aabb: Bounding box of the scene.
density_fn: Function that takes a tensor of points and returns a tensor of densities.
grid_resolution: Resolution of the density grid.
num_samples: Number of samples per ray.
"""
super().__init__(num_samples=num_samples)
self.aabb = aabb
self.density_fn = density_fn
@ -398,7 +397,7 @@ class VolumetricSampler(Sampler):
ray_bundle: RayBundle,
num_samples: Optional[int] = None,
near_plane: float = 0.0,
remove_occluded_samples: bool = True,
remove_occluded_samples: bool = False,
) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]:
"""Generate ray samples in a bounding box.