From 32943caf5b5c6cd50ac24eeede5da79102075c11 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Sun, 18 Sep 2022 13:21:50 -0700 Subject: [PATCH] Update to nerfacc 0.0.8 (#471) --- nerfactory/models/instant_ngp.py | 20 ++++++------- nerfactory/models/modules/ray_sampler.py | 37 ++++++++++++++++-------- pyproject.toml | 2 +- 3 files changed, 35 insertions(+), 24 deletions(-) diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 26503b714..d85cadf37 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -62,9 +62,6 @@ class InstantNGPModelConfig(cfg.ModelConfig): """Instant NGP doesn't use a collider.""" field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" - # TODO(ethan): remove the density field specified here - enable_density_field: bool = False - """Whether to create a density field to filter samples.""" num_samples: int = 1024 """Number of samples in field evaluation.""" grid_resolution: int = 128 @@ -129,7 +126,7 @@ class NGPModel(Model): TrainingCallback( where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], update_every_num_iters=1, - func=self.sampler.occ_field.every_n_step, # will take in step + func=self.sampler.occ_field.every_n_step, ), ] @@ -142,11 +139,10 @@ class NGPModel(Model): def get_outputs(self, ray_bundle: RayBundle): assert self.field is not None - - n_rays = ray_bundle.origins.shape[0] + num_rays = len(ray_bundle) with torch.no_grad(): - ray_samples, packed_info = self.sampler( + ray_samples, packed_info, ray_indices = self.sampler( ray_bundle=ray_bundle, near_plane=self.config.near_plane, ) @@ -154,7 +150,7 @@ class NGPModel(Model): field_outputs = self.field(ray_samples) # accumulation - weights, ray_indices = nerfacc.volumetric_rendering_weights( + weights = nerfacc.volumetric_rendering_weights( packed_info=packed_info, sigmas=field_outputs[FieldHeadNames.DENSITY], frustum_starts=ray_samples.frustums.starts, @@ -165,10 +161,12 @@ class NGPModel(Model): rgb=field_outputs[FieldHeadNames.RGB], weights=weights, ray_indices=ray_indices, - num_rays=n_rays, + num_rays=num_rays, ) - depth = self.renderer_depth(weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=n_rays) - accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=n_rays) + depth = self.renderer_depth( + weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays + ) + accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays) alive_ray_mask = accumulation.squeeze(-1) > 0 outputs = { diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 51a80e914..300cc8b29 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -502,7 +502,7 @@ class VolumetricSampler(Sampler): num_samples: Optional[int] = None, near_plane: float = 0.0, remove_occluded_samples: bool = True, - ) -> Tuple[RaySamples, TensorType["total_samples", 3]]: + ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]: """Generate ray samples in a bounding box. Args: @@ -512,9 +512,10 @@ class VolumetricSampler(Sampler): remove_occluded_samples: Whether to remove occluded samples. (requires evaluating the density function) Returns: - a tuple of (ray_samples, packed_info) + a tuple of (ray_samples, packed_info, ray_indices) The ray_samples are packed, only storing the valid samples. The packed_info contains all the information to recover packed samples into unpacked mode for rendering. + The ray_indices contains the indices of the rays that each sample belongs to. """ if self.density_field is not None: @@ -528,8 +529,12 @@ class VolumetricSampler(Sampler): rays_o = ray_bundle.origins.contiguous() rays_d = ray_bundle.directions.contiguous() + if ray_bundle.camera_indices is not None: + camera_indices = ray_bundle.camera_indices.contiguous() + else: + camera_indices = None - packed_info, origins, dirs, starts, ends = nerfacc.volumetric_marching( + packed_info, starts, ends = nerfacc.volumetric_marching( rays_o=rays_o, rays_d=rays_d, aabb=self.aabb, @@ -540,18 +545,25 @@ class VolumetricSampler(Sampler): stratified=self.training, ) + ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long() + origins = rays_o[ray_indices] + dirs = rays_d[ray_indices] + if camera_indices is not None: + camera_indices = camera_indices[ray_indices] + if remove_occluded_samples: with torch.no_grad(): positions = origins + dirs * (starts + ends) / 2.0 densities = self.density_fn(positions) - packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps( - packed_info, - densities, - starts, - ends, - origins, - dirs, - ) + if camera_indices is not None: + packed_info, starts, ends, origins, dirs, camera_indices = nerfacc.volumetric_rendering_steps( + packed_info, densities, starts, ends, origins, dirs, camera_indices + ) + else: + packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps( + packed_info, densities, starts, ends, origins, dirs + ) + ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long() zeros = torch.zeros_like(origins[:, :1]) ray_samples = RaySamples( @@ -562,5 +574,6 @@ class VolumetricSampler(Sampler): ends=ends, pixel_area=zeros, ), + camera_indices=camera_indices, ) - return ray_samples, packed_info + return ray_samples, packed_info, ray_indices diff --git a/pyproject.toml b/pyproject.toml index 03ccb1b84..f85bc3665 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "mediapy==1.1.0", "msgpack==1.0.4", "msgpack_numpy==0.4.8", - "nerfacc==0.0.6", + "nerfacc==0.0.8", "numpy-quaternion", "opencv-python==4.6.0.66", "plotly==5.7.0",