From de6af6c4e519a93d781fdf58dc97bedfb234a7ff Mon Sep 17 00:00:00 2001 From: Saurav Maheshkar Date: Tue, 14 Mar 2023 17:25:48 +0000 Subject: [PATCH] feat(fields): create common `get_normalized_directions` fn (#1599) * feat(fields): create common get_normalized_directions fn * style: run black * refactor: rename func and update docstring --- nerfstudio/fields/base_field.py | 9 +++++++++ nerfstudio/fields/instant_ngp_field.py | 13 ++----------- nerfstudio/fields/nerfacto_field.py | 13 ++----------- nerfstudio/fields/nerfplayer_nerfacto_field.py | 13 ++----------- nerfstudio/fields/nerfplayer_ngp_field.py | 5 ++--- 5 files changed, 17 insertions(+), 36 deletions(-) diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 3bcc51d56..d8e4713d1 100644 --- a/nerfstudio/fields/base_field.py +++ b/nerfstudio/fields/base_field.py @@ -126,3 +126,12 @@ class Field(nn.Module): normals = self.get_normals() field_outputs[FieldHeadNames.NORMALS] = normals # type: ignore return field_outputs + + +def shift_directions_for_tcnn(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]: + """Shift directions from [-1, 1] to [0, 1] + + Args: + directions: batch of directions + """ + return (directions + 1.0) / 2.0 diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index b0f4e97e4..d1b3e2fa7 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/nerfstudio/fields/instant_ngp_field.py @@ -30,7 +30,7 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.activations import trunc_exp from nerfstudio.field_components.embedding import Embedding from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.fields.base_field import Field +from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn try: import tinycudann as tcnn @@ -39,15 +39,6 @@ except ImportError: pass -def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]: - """SH encoding must be in the range [0, 1] - - Args: - directions: batch of directions - """ - return (directions + 1.0) / 2.0 - - class TCNNInstantNGPField(Field): """TCNN implementation of the Instant-NGP field. @@ -159,7 +150,7 @@ class TCNNInstantNGPField(Field): def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: - directions = get_normalized_directions(ray_samples.frustums.directions) + directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 830962edf..de0fd7c7d 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -46,7 +46,7 @@ from nerfstudio.field_components.spatial_distortions import ( SceneContraction, SpatialDistortion, ) -from nerfstudio.fields.base_field import Field +from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn try: import tinycudann as tcnn @@ -55,15 +55,6 @@ except ImportError: pass -def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]: - """SH encoding must be in the range [0, 1] - - Args: - directions: batch of directions - """ - return (directions + 1.0) / 2.0 - - class TCNNNerfactoField(Field): """Compound Field that uses TCNN @@ -267,7 +258,7 @@ class TCNNNerfactoField(Field): if ray_samples.camera_indices is None: raise AttributeError("Camera indices are not provided.") camera_indices = ray_samples.camera_indices.squeeze() - directions = get_normalized_directions(ray_samples.frustums.directions) + directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py index 0b186464f..cb21013e6 100644 --- a/nerfstudio/fields/nerfplayer_nerfacto_field.py +++ b/nerfstudio/fields/nerfplayer_nerfacto_field.py @@ -38,7 +38,7 @@ from nerfstudio.field_components.field_heads import ( ) from nerfstudio.field_components.spatial_distortions import SpatialDistortion from nerfstudio.field_components.temporal_grid import TemporalGridEncoder -from nerfstudio.fields.base_field import Field +from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn try: import tinycudann as tcnn @@ -47,15 +47,6 @@ except ImportError: pass -def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]: - """SH encoding must be in the range [0, 1] - - Args: - directions: batch of directions - """ - return (directions + 1.0) / 2.0 - - class TemporalHashMLPDensityField(Field): """A lightweight temporal density field module. @@ -347,7 +338,7 @@ class NerfplayerNerfactoField(Field): if ray_samples.camera_indices is None: raise AttributeError("Camera indices are not provided.") camera_indices = ray_samples.camera_indices.squeeze() - directions = get_normalized_directions(ray_samples.frustums.directions) + directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py index 829ff0d33..2952a4369 100644 --- a/nerfstudio/fields/nerfplayer_ngp_field.py +++ b/nerfstudio/fields/nerfplayer_ngp_field.py @@ -30,8 +30,7 @@ from nerfstudio.field_components.activations import trunc_exp from nerfstudio.field_components.embedding import Embedding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.temporal_grid import TemporalGridEncoder -from nerfstudio.fields.base_field import Field -from nerfstudio.fields.instant_ngp_field import get_normalized_directions +from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn try: import tinycudann as tcnn @@ -164,7 +163,7 @@ class NerfplayerNGPField(Field): def get_outputs( self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None ) -> Dict[FieldHeadNames, TensorType]: - directions = get_normalized_directions(ray_samples.frustums.directions) + directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) if self.direction_encoding is not None: