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

feat(fields): create common get_normalized_directions fn (#1599)

Browse Source 
* feat(fields): create common get_normalized_directions fn

* style: run black

* refactor: rename func and update docstring
This commit is contained in:
Saurav Maheshkar 2023-03-14 17:25:48 +00:00 committed by GitHub
parent 630babacd6
commit de6af6c4e5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 17 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
nerfstudio/fields/base_field.py
View File

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

13
nerfstudio/fields/instant_ngp_field.py
View File

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

13
nerfstudio/fields/nerfacto_field.py
View File

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

13
nerfstudio/fields/nerfplayer_nerfacto_field.py
View File

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

5
nerfstudio/fields/nerfplayer_ngp_field.py
View File

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