From 0cb98fbdd884b1d465e59355e54d429c573d777c Mon Sep 17 00:00:00 2001 From: Saurav Maheshkar Date: Fri, 14 Jul 2023 05:44:07 +0530 Subject: [PATCH] refactor: `get_normalized_directions` (#2216) * refactor: normalized_directions * fix: directions shape --- nerfstudio/fields/base_field.py | 4 ++-- nerfstudio/fields/generfacto_field.py | 20 +++++--------------- nerfstudio/fields/nerfacto_field.py | 4 ++-- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 068d4e7f1..b0215aae6 100644 --- a/nerfstudio/fields/base_field.py +++ b/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 diff --git a/nerfstudio/fields/generfacto_field.py b/nerfstudio/fields/generfacto_field.py index d72f5c263..ee31bcb0b 100644 --- a/nerfstudio/fields/generfacto_field.py +++ b/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): diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 1c35b7121..3ba9efb82 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/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)