refactor: get_normalized_directions (#2216)

* refactor: normalized_directions

* fix: directions shape

nerfstudio/fields/base_field.py

@@ -133,8 +133,8 @@ class Field(nn.Module):
         return field_outputs
 
 
-def shift_directions_for_tcnn(directions: Float[Tensor, "*bs 3"]) -> Float[Tensor, "*bs 3"]:
-    """Shift directions from [-1, 1] to [0, 1]
+def get_normalized_directions(directions: Float[Tensor, "*bs 3"]) -> Float[Tensor, "*bs 3"]:
+    """SH encoding must be in the range [0, 1]
 
     Args:
         directions: batch of directions

nerfstudio/fields/generfacto_field.py

@@ -17,30 +17,20 @@ Field for Generfacto model
 """
 
 
-from typing import Optional, Tuple, Dict, Literal
+from typing import Dict, Literal, Optional, Tuple
 
 import numpy as np
 import torch
-from torch.nn.parameter import Parameter
-from jaxtyping import Float
 from torch import Tensor, nn
+from torch.nn.parameter import Parameter
 
-from nerfstudio.field_components.mlp import MLP
 from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
-from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.encodings import HashEncoding, SHEncoding
-from nerfstudio.fields.base_field import Field
-
-
-def get_normalized_directions(directions: Float[Tensor, "bs 3"]):
-    """SH encoding must be in the range [0, 1]
-
-    Args:
-        directions: batch of directions
-    """
-    return (directions + 1.0) / 2.0
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.mlp import MLP
+from nerfstudio.fields.base_field import Field, get_normalized_directions
 
 
 class GenerfactoField(Field):

nerfstudio/fields/nerfacto_field.py

@@ -37,7 +37,7 @@ from nerfstudio.field_components.field_heads import (
 )
 from nerfstudio.field_components.mlp import MLP
 from nerfstudio.field_components.spatial_distortions import SpatialDistortion
-from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
+from nerfstudio.fields.base_field import Field, get_normalized_directions
 
 
 class NerfactoField(Field):
@@ -232,7 +232,7 @@ class NerfactoField(Field):
         if ray_samples.camera_indices is None:
             raise AttributeError("Camera indices are not provided.")
         camera_indices = ray_samples.camera_indices.squeeze()
-        directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
+        directions = get_normalized_directions(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
         d = self.direction_encoding(directions_flat)