# Copyright 2023 The University of York. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Modified by Mark Boss
"""RENI field"""
import contextlib
from dataclasses import dataclass
from typing import Dict, Literal, Optional
import torch
from einops.layers.torch import Rearrange
from jaxtyping import Float
from torch import Tensor, nn
from spar3d.models.network import get_activation_module, trunc_exp
from spar3d.models.utils import BaseModule
from .components.film_siren import FiLMSiren
from .components.siren import Siren
from .components.transformer_decoder import Decoder
from .components.vn_layers import VNInvariant, VNLinear
# from nerfstudio.cameras.rays import RaySamples
def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
"""Computes the expected value of sin(y) where y ~ N(x_means, x_vars)
Args:
x_means: Mean values.
x_vars: Variance of values.
Returns:
torch.Tensor: The expected value of sin.
"""
return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
class NeRFEncoding(torch.nn.Module):
"""Multi-scale sinousoidal encodings. Support ``integrated positional encodings`` if covariances are provided.
Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp.
Args:
in_dim: Input dimension of tensor
num_frequencies: Number of encoded frequencies per axis
min_freq_exp: Minimum frequency exponent
max_freq_exp: Maximum frequency exponent
include_input: Append the input coordinate to the encoding
"""
def __init__(
self,
in_dim: int,
num_frequencies: int,
min_freq_exp: float,
max_freq_exp: float,
include_input: bool = False,
off_axis: bool = False,
) -> None:
super().__init__()
self.in_dim = in_dim
self.num_frequencies = num_frequencies
self.min_freq = min_freq_exp
self.max_freq = max_freq_exp
self.include_input = include_input
self.off_axis = off_axis
self.P = torch.tensor(
[
[0.8506508, 0, 0.5257311],
[0.809017, 0.5, 0.309017],
[0.5257311, 0.8506508, 0],
[1, 0, 0],
[0.809017, 0.5, -0.309017],
[0.8506508, 0, -0.5257311],
[0.309017, 0.809017, -0.5],
[0, 0.5257311, -0.8506508],
[0.5, 0.309017, -0.809017],
[0, 1, 0],
[-0.5257311, 0.8506508, 0],
[-0.309017, 0.809017, -0.5],
[0, 0.5257311, 0.8506508],
[-0.309017, 0.809017, 0.5],
[0.309017, 0.809017, 0.5],
[0.5, 0.309017, 0.809017],
[0.5, -0.309017, 0.809017],
[0, 0, 1],
[-0.5, 0.309017, 0.809017],
[-0.809017, 0.5, 0.309017],
[-0.809017, 0.5, -0.309017],
]
).T
def get_out_dim(self) -> int:
if self.in_dim is None:
raise ValueError("Input dimension has not been set")
out_dim = self.in_dim * self.num_frequencies * 2
if self.off_axis:
out_dim = self.P.shape[1] * self.num_frequencies * 2
if self.include_input:
out_dim += self.in_dim
return out_dim
def forward(
self,
in_tensor: Float[Tensor, "*b input_dim"],
covs: Optional[Float[Tensor, "*b input_dim input_dim"]] = None,
) -> Float[Tensor, "*b output_dim"]:
"""Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed
in mip-NeRF.
Args:
in_tensor: For best performance, the input tensor should be between 0 and 1.
covs: Covariances of input points.
Returns:
Output values will be between -1 and 1
"""
# TODO check scaling here but just comment it for now
# in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi]
freqs = 2 ** torch.linspace(
self.min_freq, self.max_freq, self.num_frequencies
).to(in_tensor.device)
# freqs = 2 ** (
# torch.sin(torch.linspace(self.min_freq, torch.pi / 2.0, self.num_frequencies)) * self.max_freq
# ).to(in_tensor.device)
# freqs = 2 ** (
# torch.linspace(self.min_freq, 1.0, self.num_frequencies).to(in_tensor.device) ** 0.2 * self.max_freq
# )
if self.off_axis:
scaled_inputs = (
torch.matmul(in_tensor, self.P.to(in_tensor.device))[..., None] * freqs
)
else:
scaled_inputs = (
in_tensor[..., None] * freqs
) # [..., "input_dim", "num_scales"]
scaled_inputs = scaled_inputs.view(
*scaled_inputs.shape[:-2], -1
) # [..., "input_dim" * "num_scales"]
if covs is None:
encoded_inputs = torch.sin(
torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1)
)
else:
input_var = (
torch.diagonal(covs, dim1=-2, dim2=-1)[..., :, None]
* freqs[None, :] ** 2
)
input_var = input_var.reshape((*input_var.shape[:-2], -1))
encoded_inputs = expected_sin(
torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1),
torch.cat(2 * [input_var], dim=-1),
)
if self.include_input:
encoded_inputs = torch.cat([encoded_inputs, in_tensor], dim=-1)
return encoded_inputs
class RENIField(BaseModule):
@dataclass
class Config(BaseModule.Config):
"""Configuration for model instantiation"""
fixed_decoder: bool = False
"""Whether to fix the decoder weights"""
equivariance: str = "SO2"
"""Type of equivariance to use: None, SO2, SO3"""
axis_of_invariance: str = "y"
"""Which axis should SO2 equivariance be invariant to: x, y, z"""
invariant_function: str = "GramMatrix"
"""Type of invariant function to use: GramMatrix, VN"""
conditioning: str = "Concat"
"""Type of conditioning to use: FiLM, Concat, Attention"""
positional_encoding: str = "NeRF"
"""Type of positional encoding to use. Currently only NeRF is supported"""
encoded_input: str = "Directions"
"""Type of input to encode: None, Directions, Conditioning, Both"""
latent_dim: int = 36
"""Dimensionality of latent code, N for a latent code size of (N x 3)"""
hidden_layers: int = 3
"""Number of hidden layers"""
hidden_features: int = 128
"""Number of hidden features"""
mapping_layers: int = 3
"""Number of mapping layers"""
mapping_features: int = 128
"""Number of mapping features"""
num_attention_heads: int = 8
"""Number of attention heads"""
num_attention_layers: int = 3
"""Number of attention layers"""
out_features: int = 3 # RGB
"""Number of output features"""
last_layer_linear: bool = False
"""Whether to use a linear layer as the last layer"""
output_activation: str = "exp"
"""Activation function for output layer: sigmoid, tanh, relu, exp, None"""
first_omega_0: float = 30.0
"""Omega_0 for first layer"""
hidden_omega_0: float = 30.0
"""Omega_0 for hidden layers"""
fixed_decoder: bool = False
"""Whether to fix the decoder weights"""
old_implementation: bool = False
"""Whether to match implementation of old RENI, when using old checkpoints"""
cfg: Config
def configure(self):
self.equivariance = self.cfg.equivariance
self.conditioning = self.cfg.conditioning
self.latent_dim = self.cfg.latent_dim
self.hidden_layers = self.cfg.hidden_layers
self.hidden_features = self.cfg.hidden_features
self.mapping_layers = self.cfg.mapping_layers
self.mapping_features = self.cfg.mapping_features
self.out_features = self.cfg.out_features
self.last_layer_linear = self.cfg.last_layer_linear
self.output_activation = self.cfg.output_activation
self.first_omega_0 = self.cfg.first_omega_0
self.hidden_omega_0 = self.cfg.hidden_omega_0
self.old_implementation = self.cfg.old_implementation
self.axis_of_invariance = ["x", "y", "z"].index(self.cfg.axis_of_invariance)
self.fixed_decoder = self.cfg.fixed_decoder
if self.cfg.invariant_function == "GramMatrix":
self.invariant_function = self.gram_matrix_invariance
else:
self.vn_proj_in = nn.Sequential(
Rearrange("... c -> ... 1 c"),
VNLinear(dim_in=1, dim_out=1, bias_epsilon=0),
)
dim_coor = 2 if self.cfg.equivariance == "SO2" else 3
self.vn_invar = VNInvariant(dim=1, dim_coor=dim_coor)
self.invariant_function = self.vn_invariance
self.network = self.setup_network()
if self.fixed_decoder:
for param in self.network.parameters():
param.requires_grad = False
if self.cfg.invariant_function == "VN":
for param in self.vn_proj_in.parameters():
param.requires_grad = False
for param in self.vn_invar.parameters():
param.requires_grad = False
@contextlib.contextmanager
def hold_decoder_fixed(self):
"""Context manager to fix the decoder weights
Example usage:
```
with instance_of_RENIField.hold_decoder_fixed():
# do stuff
```
"""
prev_state_network = {
name: p.requires_grad for name, p in self.network.named_parameters()
}
for param in self.network.parameters():
param.requires_grad = False
if self.cfg.invariant_function == "VN":
prev_state_proj_in = {
k: p.requires_grad for k, p in self.vn_proj_in.named_parameters()
}
prev_state_invar = {
k: p.requires_grad for k, p in self.vn_invar.named_parameters()
}
for param in self.vn_proj_in.parameters():
param.requires_grad = False
for param in self.vn_invar.parameters():
param.requires_grad = False
prev_decoder_state = self.fixed_decoder
self.fixed_decoder = True
try:
yield
finally:
# Restore the previous requires_grad state
for name, param in self.network.named_parameters():
param.requires_grad = prev_state_network[name]
if self.cfg.invariant_function == "VN":
for name, param in self.vn_proj_in.named_parameters():
param.requires_grad_(prev_state_proj_in[name])
for name, param in self.vn_invar.named_parameters():
param.requires_grad_(prev_state_invar[name])
self.fixed_decoder = prev_decoder_state
def vn_invariance(
self,
Z: Float[Tensor, "B latent_dim 3"],
D: Float[Tensor, "B num_rays 3"],
equivariance: Literal["None", "SO2", "SO3"] = "SO2",
axis_of_invariance: int = 1,
):
"""Generates a batched invariant representation from latent code Z and direction coordinates D.
Args:
Z: [B, latent_dim, 3] - Latent code.
D: [B num_rays, 3] - Direction coordinates.
equivariance: The type of equivariance to use. Options are 'None', 'SO2', 'SO3'.
axis_of_invariance: The axis of rotation invariance. Should be 0 (x-axis), 1 (y-axis), or 2 (z-axis).
Returns:
Tuple[Tensor, Tensor]: directional_input, conditioning_input
"""
assert 0 <= axis_of_invariance < 3, "axis_of_invariance should be 0, 1, or 2."
other_axes = [i for i in range(3) if i != axis_of_invariance]
B, latent_dim, _ = Z.shape
_, num_rays, _ = D.shape
if equivariance == "None":
# get inner product between latent code and direction coordinates
innerprod = torch.sum(
Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
) # [B, num_rays, latent_dim]
z_input = (
Z.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, latent_dim * 3)
) # [B, num_rays, latent_dim * 3]
return innerprod, z_input
if equivariance == "SO2":
z_other = torch.stack(
(Z[..., other_axes[0]], Z[..., other_axes[1]]), -1
) # [B, latent_dim, 2]
d_other = torch.stack(
(D[..., other_axes[0]], D[..., other_axes[1]]), -1
).unsqueeze(2) # [B, num_rays, 1, 2]
d_other = d_other.expand(
B, num_rays, latent_dim, 2
) # [B, num_rays, latent_dim, 2]
z_other_emb = self.vn_proj_in(z_other) # [B, latent_dim, 1, 2]
z_other_invar = self.vn_invar(z_other_emb) # [B, latent_dim, 2]
# Get invariant component of Z along the axis of invariance
z_invar = Z[..., axis_of_invariance].unsqueeze(-1) # [B, latent_dim, 1]
# Innerproduct between projection of Z and D on the plane orthogonal to the axis of invariance.
# This encodes the rotational information. This is rotation-equivariant to rotations of either Z
# or D and is invariant to rotations of both Z and D.
innerprod = (z_other.unsqueeze(1) * d_other).sum(
dim=-1
) # [B, num_rays, latent_dim]
# Compute norm along the axes orthogonal to the axis of invariance
d_other_norm = torch.sqrt(
D[..., other_axes[0]] ** 2 + D[..., other_axes[1]] ** 2
).unsqueeze(-1) # [B num_rays, 1]
# Get invariant component of D along the axis of invariance
d_invar = D[..., axis_of_invariance].unsqueeze(-1) # [B, num_rays, 1]
directional_input = torch.cat(
(innerprod, d_invar, d_other_norm), -1
) # [B, num_rays, latent_dim + 2]
conditioning_input = (
torch.cat((z_other_invar, z_invar), dim=-1)
.flatten(1)
.unsqueeze(1)
.expand(B, num_rays, latent_dim * 3)
) # [B, num_rays, latent_dim * 3]
return directional_input, conditioning_input
if equivariance == "SO3":
z = self.vn_proj_in(Z) # [B, latent_dim, 1, 3]
z_invar = self.vn_invar(z) # [B, latent_dim, 3]
conditioning_input = (
z_invar.flatten(1).unsqueeze(1).expand(B, num_rays, latent_dim)
) # [B, num_rays, latent_dim * 3]
# D [B, num_rays, 3] -> [B, num_rays, 1, 3]
# Z [B, latent_dim, 3] -> [B, 1, latent_dim, 3]
innerprod = torch.sum(
Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
) # [B, num_rays, latent_dim]
return innerprod, conditioning_input
def gram_matrix_invariance(
self,
Z: Float[Tensor, "B latent_dim 3"],
D: Float[Tensor, "B num_rays 3"],
equivariance: Literal["None", "SO2", "SO3"] = "SO2",
axis_of_invariance: int = 1,
):
"""Generates an invariant representation from latent code Z and direction coordinates D.
Args:
Z (torch.Tensor): Latent code (B x latent_dim x 3)
D (torch.Tensor): Direction coordinates (B x num_rays x 3)
equivariance (str): Type of equivariance to use. Options are 'none', 'SO2', and 'SO3'
axis_of_invariance (int): The axis of rotation invariance. Should be 0 (x-axis), 1 (y-axis), or 2 (z-axis).
Default is 1 (y-axis).
Returns:
torch.Tensor: Invariant representation
"""
assert 0 <= axis_of_invariance < 3, "axis_of_invariance should be 0, 1, or 2."
other_axes = [i for i in range(3) if i != axis_of_invariance]
B, latent_dim, _ = Z.shape
_, num_rays, _ = D.shape
if equivariance == "None":
# get inner product between latent code and direction coordinates
innerprod = torch.sum(
Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
) # [B, num_rays, latent_dim]
z_input = (
Z.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, latent_dim * 3)
) # [B, num_rays, latent_dim * 3]
return innerprod, z_input
if equivariance == "SO2":
# Select components along axes orthogonal to the axis of invariance
z_other = torch.stack(
(Z[..., other_axes[0]], Z[..., other_axes[1]]), -1
) # [B, latent_dim, 2]
d_other = torch.stack(
(D[..., other_axes[0]], D[..., other_axes[1]]), -1
).unsqueeze(2) # [B, num_rays, 1, 2]
d_other = d_other.expand(
B, num_rays, latent_dim, 2
) # size becomes [B, num_rays, latent_dim, 2]
# Invariant representation of Z, gram matrix G=Z*Z' is size num_rays x latent_dim x latent_dim
G = torch.bmm(z_other, torch.transpose(z_other, 1, 2))
# Flatten G to be size B x latent_dim^2
z_other_invar = G.flatten(start_dim=1)
# Get invariant component of Z along the axis of invariance
z_invar = Z[..., axis_of_invariance] # [B, latent_dim]
# Innerprod is size num_rays x latent_dim
innerprod = (z_other.unsqueeze(1) * d_other).sum(
dim=-1
) # [B, num_rays, latent_dim]
# Compute norm along the axes orthogonal to the axis of invariance
d_other_norm = torch.sqrt(
D[..., other_axes[0]] ** 2 + D[..., other_axes[1]] ** 2
).unsqueeze(-1) # [B, num_rays, 1]
# Get invariant component of D along the axis of invariance
d_invar = D[..., axis_of_invariance].unsqueeze(-1) # [B, num_rays, 1]
if not self.old_implementation:
directional_input = torch.cat(
(innerprod, d_invar, d_other_norm), -1
) # [B, num_rays, latent_dim + 2]
conditioning_input = (
torch.cat((z_other_invar, z_invar), -1)
.unsqueeze(1)
.expand(B, num_rays, latent_dim * 3)
) # [B, num_rays, latent_dim^2 + latent_dim]
else:
# this is matching the previous implementation of RENI, needed if using old checkpoints
z_other_invar = z_other_invar.unsqueeze(1).expand(B, num_rays, -1)
z_invar = z_invar.unsqueeze(1).expand(B, num_rays, -1)
return torch.cat(
(innerprod, z_other_invar, d_other_norm, z_invar, d_invar), 1
)
return directional_input, conditioning_input
if equivariance == "SO3":
G = Z @ torch.transpose(Z, 1, 2) # [B, latent_dim, latent_dim]
innerprod = torch.sum(
Z.unsqueeze(1) * D.unsqueeze(2), dim=-1
) # [B, num_rays, latent_dim]
z_invar = (
G.flatten(start_dim=1).unsqueeze(1).expand(B, num_rays, -1)
) # [B, num_rays, latent_dim^2]
return innerprod, z_invar
def setup_network(self):
"""Sets up the network architecture"""
base_input_dims = {
"VN": {
"None": {
"direction": self.latent_dim,
"conditioning": self.latent_dim * 3,
},
"SO2": {
"direction": self.latent_dim + 2,
"conditioning": self.latent_dim * 3,
},
"SO3": {
"direction": self.latent_dim,
"conditioning": self.latent_dim * 3,
},
},
"GramMatrix": {
"None": {
"direction": self.latent_dim,
"conditioning": self.latent_dim * 3,
},
"SO2": {
"direction": self.latent_dim + 2,
"conditioning": self.latent_dim**2 + self.latent_dim,
},
"SO3": {
"direction": self.latent_dim,
"conditioning": self.latent_dim**2,
},
},
}
# Extract the necessary input dimensions
input_types = ["direction", "conditioning"]
input_dims = {
key: base_input_dims[self.cfg.invariant_function][self.cfg.equivariance][
key
]
for key in input_types
}
# Helper function to create NeRF encoding
def create_nerf_encoding(in_dim):
return NeRFEncoding(
in_dim=in_dim,
num_frequencies=2,
min_freq_exp=0.0,
max_freq_exp=2.0,
include_input=True,
)
# Dictionary-based encoding setup
encoding_setup = {
"None": [],
"Conditioning": ["conditioning"],
"Directions": ["direction"],
"Both": ["direction", "conditioning"],
}
# Setting up the required encodings
for input_type in encoding_setup.get(self.cfg.encoded_input, []):
# create self.{input_type}_encoding and update input_dims
setattr(
self,
f"{input_type}_encoding",
create_nerf_encoding(input_dims[input_type]),
)
input_dims[input_type] = getattr(
self, f"{input_type}_encoding"
).get_out_dim()
output_activation = get_activation_module(self.cfg.output_activation)
network = None
if self.conditioning == "Concat":
network = Siren(
in_dim=input_dims["direction"] + input_dims["conditioning"],
hidden_layers=self.hidden_layers,
hidden_features=self.hidden_features,
out_dim=self.out_features,
outermost_linear=self.last_layer_linear,
first_omega_0=self.first_omega_0,
hidden_omega_0=self.hidden_omega_0,
out_activation=output_activation,
)
elif self.conditioning == "FiLM":
network = FiLMSiren(
in_dim=input_dims["direction"],
hidden_layers=self.hidden_layers,
hidden_features=self.hidden_features,
mapping_network_in_dim=input_dims["conditioning"],
mapping_network_layers=self.mapping_layers,
mapping_network_features=self.mapping_features,
out_dim=self.out_features,
outermost_linear=True,
out_activation=output_activation,
)
elif self.conditioning == "Attention":
# transformer where K, V is from conditioning input and Q is from pos encoded directional input
network = Decoder(
in_dim=input_dims["direction"],
conditioning_input_dim=input_dims["conditioning"],
hidden_features=self.cfg.hidden_features,
num_heads=self.cfg.num_attention_heads,
num_layers=self.cfg.num_attention_layers,
out_activation=output_activation,
)
assert network is not None, "unknown conditioning type"
return network
def apply_positional_encoding(self, directional_input, conditioning_input):
# conditioning on just invariant directional input
if self.cfg.encoded_input == "Conditioning":
conditioning_input = self.conditioning_encoding(
conditioning_input
) # [num_rays, embedding_dim]
elif self.cfg.encoded_input == "Directions":
directional_input = self.direction_encoding(
directional_input
) # [num_rays, embedding_dim]
elif self.cfg.encoded_input == "Both":
directional_input = self.direction_encoding(directional_input)
conditioning_input = self.conditioning_encoding(conditioning_input)
return directional_input, conditioning_input
def get_outputs(
self,
rays_d: Float[Tensor, "batch num_rays 3"], # type: ignore
latent_codes: Float[Tensor, "batch_size latent_dim 3"], # type: ignore
rotation: Optional[Float[Tensor, "batch_size 3 3"]] = None, # type: ignore
scale: Optional[Float[Tensor, "batch_size"]] = None, # type: ignore
) -> Dict[str, Tensor]:
"""Returns the outputs of the field.
Args:
ray_samples: [batch_size num_rays 3]
latent_codes: [batch_size, latent_dim, 3]
rotation: [batch_size, 3, 3]
scale: [batch_size]
"""
if rotation is not None:
if len(rotation.shape) == 3: # [batch_size, 3, 3]
# Expand latent_codes to match [batch_size, latent_dim, 3]
latent_codes = torch.einsum(
"bik,blk->bli",
rotation,
latent_codes,
)
else:
raise NotImplementedError(
"Unsupported rotation shape. Expected [batch_size, 3, 3]."
)
B, num_rays, _ = rays_d.shape
_, latent_dim, _ = latent_codes.shape
if not self.old_implementation:
directional_input, conditioning_input = self.invariant_function(
latent_codes,
rays_d,
equivariance=self.equivariance,
axis_of_invariance=self.axis_of_invariance,
) # [B, num_rays, 3]
if self.cfg.positional_encoding == "NeRF":
directional_input, conditioning_input = self.apply_positional_encoding(
directional_input, conditioning_input
)
if self.conditioning == "Concat":
model_outputs = self.network(
torch.cat((directional_input, conditioning_input), dim=-1).reshape(
B * num_rays, -1
)
).view(B, num_rays, 3) # returns -> [B num_rays, 3]
elif self.conditioning == "FiLM":
model_outputs = self.network(
directional_input.reshape(B * num_rays, -1),
conditioning_input.reshape(B * num_rays, -1),
).view(B, num_rays, 3) # returns -> [B num_rays, 3]
elif self.conditioning == "Attention":
model_outputs = self.network(
directional_input.reshape(B * num_rays, -1),
conditioning_input.reshape(B * num_rays, -1),
).view(B, num_rays, 3) # returns -> [B num_rays, 3]
else:
# in the old implementation directions were sampled with y-up not z-up so need to swap y and z in directions
directions = torch.stack(
(rays_d[..., 0], rays_d[..., 2], rays_d[..., 1]), -1
)
model_input = self.invariant_function(
latent_codes,
directions,
equivariance=self.equivariance,
axis_of_invariance=self.axis_of_invariance,
) # [B, num_rays, 3]
model_outputs = self.network(model_input.view(B * num_rays, -1)).view(
B, num_rays, 3
)
outputs = {}
if scale is not None:
scale = trunc_exp(scale) # [num_rays] exp to ensure positive
model_outputs = model_outputs * scale.view(-1, 1, 1) # [num_rays, 3]
outputs["rgb"] = model_outputs
return outputs
def forward(
self,
rays_d: Float[Tensor, "batch num_rays 3"], # type: ignore
latent_codes: Float[Tensor, "batch_size latent_dim 3"], # type: ignore
rotation: Optional[Float[Tensor, "batch_size 3 3"]] = None, # type: ignore
scale: Optional[Float[Tensor, "batch_size"]] = None, # type: ignore
) -> Dict[str, Tensor]:
"""Evaluates spherical field for a given ray bundle and rotation.
Args:
ray_samples: [B num_rays 3]
latent_codes: [B, num_rays, latent_dim, 3]
rotation: [batch_size, 3, 3]
scale: [batch_size]
Returns:
Dict[str, Tensor]: A dictionary containing the outputs of the field.
"""
return self.get_outputs(
rays_d=rays_d,
latent_codes=latent_codes,
rotation=rotation,
scale=scale,
)