Merge remote-tracking branch 'origin/main' into generative

nerfstudio/fields/nerfplayer_nerfacto_field.py

@@ -0,0 +1,415 @@
+# Copyright 2022 The Nerfstudio Team. 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.
+
+"""
+Field implementations for NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backbone
+"""
+
+
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.nn.parameter import Parameter
+from torchtyping import TensorType
+
+from nerfstudio.cameras.rays import Frustums, RaySamples
+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,
+    PredNormalsFieldHead,
+    SemanticFieldHead,
+    TransientDensityFieldHead,
+    TransientRGBFieldHead,
+    UncertaintyFieldHead,
+)
+from nerfstudio.field_components.spatial_distortions import SpatialDistortion
+from nerfstudio.field_components.temporal_grid import TemporalGridEncoder
+from nerfstudio.fields.base_field import Field
+
+try:
+    import tinycudann as tcnn
+except ImportError:
+    # tinycudann module doesn't exist
+    pass
+
+
+def get_normalized_directions(directions: 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.
+
+    Args:
+        aabb: Parameters of scene aabb bounds
+        temporal_dim: Hashing grid parameter. A higher temporal dim means a higher temporal frequency.
+        num_layers: Number of hidden layers
+        hidden_dim: Dimension of hidden layers
+        spatial_distortion: Spatial distortion module
+        num_levels: Hashing grid parameter. Used for initialize TemporalGridEncoder class.
+        max_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class.
+        base_res: Hashing grid parameter. Used for initialize TemporalGridEncoder class.
+        log2_hashmap_size: Hashing grid parameter. Used for initialize TemporalGridEncoder class.
+        features_per_level: Hashing grid parameter. Used for initialize TemporalGridEncoder class.
+    """
+
+    def __init__(
+        self,
+        aabb,
+        temporal_dim: int = 64,
+        num_layers: int = 2,
+        hidden_dim: int = 64,
+        spatial_distortion: Optional[SpatialDistortion] = None,
+        num_levels=8,
+        max_res=1024,
+        base_res=16,
+        log2_hashmap_size=18,
+        features_per_level=2,
+    ) -> None:
+        super().__init__()
+        # from .temporal_grid import test; test() # DEBUG
+        self.aabb = Parameter(aabb, requires_grad=False)
+        self.spatial_distortion = spatial_distortion
+        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
+
+        self.encoding = TemporalGridEncoder(
+            input_dim=3,
+            temporal_dim=temporal_dim,
+            num_levels=num_levels,
+            level_dim=features_per_level,
+            per_level_scale=growth_factor,
+            base_resolution=base_res,
+            log2_hashmap_size=log2_hashmap_size,
+        )
+        self.linear = tcnn.Network(
+            n_input_dims=num_levels * features_per_level,
+            n_output_dims=1,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": hidden_dim,
+                "n_hidden_layers": num_layers - 1,
+            },
+        )
+
+    # pylint: disable=arguments-differ
+    def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs", 1]) -> TensorType["bs":..., 1]:
+        """Returns only the density. Used primarily with the density grid.
+
+        Args:
+            positions: the origin of the samples/frustums
+            times: the time of rays
+        """
+        if len(positions.shape) == 3 and len(times.shape) == 2:
+            # position is [ray, sample, 3]; times is [ray, 1]
+            times = times[:, None]  # RaySamples can handle the shape
+        # Need to figure out a better way to descibe positions with a ray.
+        ray_samples = RaySamples(
+            frustums=Frustums(
+                origins=positions,
+                directions=torch.ones_like(positions),
+                starts=torch.zeros_like(positions[..., :1]),
+                ends=torch.zeros_like(positions[..., :1]),
+                pixel_area=torch.ones_like(positions[..., :1]),
+            ),
+            times=times,
+        )
+        density, _ = self.get_density(ray_samples)
+        return density
+
+    def get_density(self, ray_samples: RaySamples):
+        if self.spatial_distortion is not None:
+            positions = self.spatial_distortion(ray_samples.frustums.get_positions())
+            positions = (positions + 2.0) / 4.0
+        else:
+            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions_flat = positions.view(-1, 3)
+        time_flat = ray_samples.times.reshape(-1, 1)
+        x = self.encoding(positions_flat, time_flat).to(positions)
+        density_before_activation = self.linear(x).view(*ray_samples.frustums.shape, -1)
+
+        # Rectifying the density with an exponential is much more stable than a ReLU or
+        # softplus, because it enables high post-activation (float32) density outputs
+        # from smaller internal (float16) parameters.
+        density = trunc_exp(density_before_activation)
+        return density, None
+
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+        return {}
+
+
+class NerfplayerNerfactoField(Field):
+    """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with nerfacto backbone.
+
+    Args:
+        aabb: parameters of scene aabb bounds
+        num_images: number of images in the dataset
+        num_layers: number of hidden layers
+        hidden_dim: dimension of hidden layers
+        geo_feat_dim: output geo feat dimensions
+        num_levels: number of levels of the hashmap for the base mlp
+        max_res: maximum resolution of the hashmap for the base mlp
+        log2_hashmap_size: size of the hashmap for the base mlp
+        num_layers_color: number of hidden layers for color network
+        num_layers_transient: number of hidden layers for transient network
+        hidden_dim_color: dimension of hidden layers for color network
+        hidden_dim_transient: dimension of hidden layers for transient network
+        appearance_embedding_dim: dimension of appearance embedding
+        transient_embedding_dim: dimension of transient embedding
+        use_transient_embedding: whether to use transient embedding
+        use_semantics: whether to use semantic segmentation
+        num_semantic_classes: number of semantic classes
+        use_pred_normals: whether to use predicted normals
+        use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
+        spatial_distortion: spatial distortion to apply to the scene
+    """
+
+    def __init__(
+        self,
+        aabb,
+        num_images: int,
+        num_layers: int = 2,
+        hidden_dim: int = 64,
+        geo_feat_dim: int = 15,
+        temporal_dim: int = 64,
+        num_levels: int = 16,
+        features_per_level: int = 2,
+        log2_hashmap_size: int = 19,
+        num_layers_color: int = 3,
+        num_layers_transient: int = 2,
+        hidden_dim_color: int = 64,
+        hidden_dim_transient: int = 64,
+        appearance_embedding_dim: int = 32,
+        transient_embedding_dim: int = 16,
+        use_transient_embedding: bool = False,
+        use_semantics: bool = False,
+        num_semantic_classes: int = 100,
+        use_pred_normals: bool = False,
+        use_average_appearance_embedding: bool = False,
+        spatial_distortion: Optional[SpatialDistortion] = None,
+    ) -> None:
+        super().__init__()
+
+        self.aabb = Parameter(aabb, requires_grad=False)
+        self.geo_feat_dim = geo_feat_dim
+
+        self.spatial_distortion = spatial_distortion
+        self.num_images = num_images
+        self.appearance_embedding_dim = appearance_embedding_dim
+        self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
+        self.use_average_appearance_embedding = use_average_appearance_embedding
+        self.use_transient_embedding = use_transient_embedding
+        self.use_semantics = use_semantics
+        self.use_pred_normals = use_pred_normals
+
+        self.direction_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 4,
+            },
+        )
+
+        self.position_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        )
+
+        self.mlp_base = TemporalGridEncoder(
+            input_dim=3,
+            temporal_dim=temporal_dim,
+            num_levels=num_levels,
+            level_dim=features_per_level,
+            log2_hashmap_size=log2_hashmap_size,
+            desired_resolution=1024 * (self.aabb.max() - self.aabb.min()),
+        )
+        self.mlp_base_decode = tcnn.Network(
+            n_input_dims=num_levels * features_per_level,
+            n_output_dims=1 + self.geo_feat_dim,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": hidden_dim,
+                "n_hidden_layers": num_layers - 1,
+            },
+        )
+
+        # transients
+        if self.use_transient_embedding:
+            self.transient_embedding_dim = transient_embedding_dim
+            self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim)
+            self.mlp_transient = tcnn.Network(
+                n_input_dims=self.geo_feat_dim + self.transient_embedding_dim,
+                n_output_dims=hidden_dim_transient,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": hidden_dim_transient,
+                    "n_hidden_layers": num_layers_transient - 1,
+                },
+            )
+            self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.n_output_dims)
+            self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.n_output_dims)
+            self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.n_output_dims)
+
+        # semantics
+        if self.use_semantics:
+            self.mlp_semantics = tcnn.Network(
+                n_input_dims=self.geo_feat_dim,
+                n_output_dims=hidden_dim_transient,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": 64,
+                    "n_hidden_layers": 1,
+                },
+            )
+            self.field_head_semantics = SemanticFieldHead(
+                in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes
+            )
+
+        # predicted normals
+        if self.use_pred_normals:
+            self.mlp_pred_normals = tcnn.Network(
+                n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
+                n_output_dims=hidden_dim_transient,
+                network_config={
+                    "otype": "FullyFusedMLP",
+                    "activation": "ReLU",
+                    "output_activation": "None",
+                    "n_neurons": 64,
+                    "n_hidden_layers": 2,
+                },
+            )
+            self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
+
+        self.mlp_head = tcnn.Network(
+            n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim + self.appearance_embedding_dim,
+            n_output_dims=3,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "Sigmoid",
+                "n_neurons": hidden_dim_color,
+                "n_hidden_layers": num_layers_color - 1,
+            },
+        )
+
+    def get_density(self, ray_samples: RaySamples):
+        """Computes and returns the densities."""
+        if self.spatial_distortion is not None:
+            positions = ray_samples.frustums.get_positions()
+            positions = self.spatial_distortion(positions)
+            positions = (positions + 2.0) / 4.0
+        else:
+            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions_flat = positions.view(-1, 3)
+        assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer"
+        time_flat = ray_samples.times.reshape(-1, 1)
+        h = self.mlp_base(positions_flat, time_flat)
+        h = self.mlp_base_decode(h).view(*ray_samples.frustums.shape, -1)
+        density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
+
+        # Rectifying the density with an exponential is much more stable than a ReLU or
+        # softplus, because it enables high post-activation (float32) density outputs
+        # from smaller internal (float16) parameters.
+        density = trunc_exp(density_before_activation.to(positions))
+        return density, base_mlp_out
+
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+        assert density_embedding is not None
+        outputs = {}
+        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_flat = directions.view(-1, 3)
+        d = self.direction_encoding(directions_flat)
+
+        outputs_shape = ray_samples.frustums.directions.shape[:-1]
+
+        # appearance
+        if self.training:
+            embedded_appearance = self.embedding_appearance(camera_indices)
+        else:
+            if self.use_average_appearance_embedding:
+                embedded_appearance = torch.ones(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                ) * self.embedding_appearance.mean(dim=0)
+            else:
+                embedded_appearance = torch.zeros(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                )
+
+        # transients
+        if self.use_transient_embedding and self.training:
+            embedded_transient = self.embedding_transient(camera_indices)
+            transient_input = torch.cat(
+                [
+                    density_embedding.view(-1, self.geo_feat_dim),
+                    embedded_transient.view(-1, self.transient_embedding_dim),
+                ],
+                dim=-1,
+            )
+            x = self.mlp_transient(transient_input).view(*outputs_shape, -1).to(directions)
+            outputs[FieldHeadNames.UNCERTAINTY] = self.field_head_transient_uncertainty(x)
+            outputs[FieldHeadNames.TRANSIENT_RGB] = self.field_head_transient_rgb(x)
+            outputs[FieldHeadNames.TRANSIENT_DENSITY] = self.field_head_transient_density(x)
+
+        # semantics
+        if self.use_semantics:
+            density_embedding_copy = density_embedding.clone().detach()
+            semantics_input = torch.cat(
+                [
+                    density_embedding_copy.view(-1, self.geo_feat_dim),
+                ],
+                dim=-1,
+            )
+            x = self.mlp_semantics(semantics_input).view(*outputs_shape, -1).to(directions)
+            outputs[FieldHeadNames.SEMANTICS] = self.field_head_semantics(x)
+
+        # predicted normals
+        if self.use_pred_normals:
+            positions = ray_samples.frustums.get_positions()
+
+            positions_flat = self.position_encoding(positions.view(-1, 3))
+            pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+
+            x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
+            outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
+
+        h = torch.cat(
+            [
+                d,
+                density_embedding.view(-1, self.geo_feat_dim),
+                embedded_appearance.view(-1, self.appearance_embedding_dim),
+            ],
+            dim=-1,
+        )
+        rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions)
+        outputs.update({FieldHeadNames.RGB: rgb})
+
+        return outputs

nerfstudio/fields/nerfplayer_ngp_field.py

@@ -0,0 +1,236 @@
+# Copyright 2022 The Nerfstudio Team. 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.
+
+"""
+NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with InstantNGP backbone
+"""
+
+
+from typing import Optional
+
+import torch
+from nerfacc import ContractionType, contract
+from torch.nn.parameter import Parameter
+from torchtyping import TensorType
+
+from nerfstudio.cameras.rays import Frustums, RaySamples
+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.field_components.temporal_grid import TemporalGridEncoder
+from nerfstudio.fields.base_field import Field
+from nerfstudio.fields.instant_ngp_field import get_normalized_directions
+
+try:
+    import tinycudann as tcnn
+except ImportError:
+    # tinycudann module doesn't exist
+    pass
+
+
+class NerfplayerNGPField(Field):
+    """NeRFPlayer (https://arxiv.org/abs/2210.15947) field with InstantNGP backbone.
+
+    Args:
+        aabb: parameters of scene aabb bounds
+        temporal_dim: the dimension of temporal axis, a higher dimension indicates a higher temporal frequency
+            please refer to the implementation of TemporalGridEncoder for more details
+        num_levels: the number of multi-resolution levels; same as InstantNGP
+        features_per_level: the dim of output feature vector for each level; same as InstantNGP
+        log2_hashmap_size: the size of the table; same as InstantNGP
+        base_resolution: base resolution for the table; same as InstantNGP
+        num_layers: number of hidden layers (occupancy decoder network after sampling)
+        hidden_dim: dimension of hidden layers (occupancy decoder network after sampling)
+        geo_feat_dim: output geo feat dimensions
+        num_layers_color: number of hidden layers for color network
+        hidden_dim_color: dimension of hidden layers for color network
+        use_appearance_embedding: whether to use appearance embedding
+        disable_viewing_dependent: if true, disable the viewing dependent effect (no viewing direction as inputs)
+            Sometimes we need to disable viewing dependent effects in a dynamic scene, because there is
+            ambiguity between being dynamic and viewing dependent effects. For example, the shadow of the camera
+            should be a dynamic effect, but may be reconstructed as viewing dependent effects.
+        num_images: number of images, requried if use_appearance_embedding is True
+        appearance_embedding_dim: dimension of appearance embedding
+        contraction_type: type of contraction
+    """
+
+    def __init__(
+        self,
+        aabb,
+        temporal_dim: int = 16,
+        num_levels: int = 16,
+        features_per_level: int = 2,
+        log2_hashmap_size: int = 19,
+        base_resolution: int = 16,
+        num_layers: int = 2,
+        hidden_dim: int = 64,
+        geo_feat_dim: int = 15,
+        num_layers_color: int = 3,
+        hidden_dim_color: int = 64,
+        use_appearance_embedding: bool = False,
+        disable_viewing_dependent: bool = False,
+        num_images: Optional[int] = None,
+        appearance_embedding_dim: int = 32,
+        contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,
+    ) -> None:
+        super().__init__()
+
+        self.aabb = Parameter(aabb, requires_grad=False)
+        self.geo_feat_dim = geo_feat_dim
+        self.contraction_type = contraction_type
+
+        self.use_appearance_embedding = use_appearance_embedding
+        if use_appearance_embedding:
+            assert num_images is not None
+            self.appearance_embedding_dim = appearance_embedding_dim
+            self.appearance_embedding = Embedding(num_images, appearance_embedding_dim)
+
+        self.direction_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 4,
+            },
+        )
+
+        self.mlp_base = TemporalGridEncoder(
+            input_dim=3,
+            temporal_dim=temporal_dim,
+            num_levels=num_levels,
+            level_dim=features_per_level,
+            base_resolution=base_resolution,
+            log2_hashmap_size=log2_hashmap_size,
+            desired_resolution=1024 * (self.aabb.max() - self.aabb.min()),
+        )
+        self.mlp_base_decode = tcnn.Network(
+            n_input_dims=num_levels * features_per_level,
+            n_output_dims=1 + self.geo_feat_dim,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": hidden_dim,
+                "n_hidden_layers": num_layers - 1,
+            },
+        )
+
+        in_dim = self.direction_encoding.n_output_dims + self.geo_feat_dim
+        if disable_viewing_dependent:
+            in_dim = self.geo_feat_dim
+            self.direction_encoding = None
+        if self.use_appearance_embedding:
+            in_dim += self.appearance_embedding_dim
+        self.mlp_head = tcnn.Network(
+            n_input_dims=in_dim,
+            n_output_dims=3,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "Sigmoid",
+                "n_neurons": hidden_dim_color,
+                "n_hidden_layers": num_layers_color - 1,
+            },
+        )
+
+    def get_density(self, ray_samples: RaySamples):
+        positions = ray_samples.frustums.get_positions()
+        positions_flat = positions.view(-1, 3)
+        positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)
+        assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer"
+        times_flat = ray_samples.times.view(-1, 1)
+
+        h = self.mlp_base(positions_flat, times_flat)
+        h = self.mlp_base_decode(h).view(*ray_samples.frustums.shape, -1)
+        density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
+
+        # Rectifying the density with an exponential is much more stable than a ReLU or
+        # softplus, because it enables high post-activation (float32) density outputs
+        # from smaller internal (float16) parameters.
+        density = trunc_exp(density_before_activation.to(positions))
+        return density, base_mlp_out
+
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+        directions = get_normalized_directions(ray_samples.frustums.directions)
+        directions_flat = directions.view(-1, 3)
+
+        if self.direction_encoding is not None:
+            d = self.direction_encoding(directions_flat)
+            if density_embedding is None:
+                positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+                h = torch.cat([d, positions.view(-1, 3)], dim=-1)
+            else:
+                h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+        else:
+            # viewing direction is disabled
+            if density_embedding is None:
+                positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+                h = positions.view(-1, 3)
+            else:
+                h = density_embedding.view(-1, self.geo_feat_dim)
+
+        if self.use_appearance_embedding:
+            if ray_samples.camera_indices is None:
+                raise AttributeError("Camera indices are not provided.")
+            camera_indices = ray_samples.camera_indices.squeeze()
+            if self.training:
+                embedded_appearance = self.appearance_embedding(camera_indices)
+            else:
+                embedded_appearance = torch.zeros(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                )
+            h = torch.cat([h, embedded_appearance.view(-1, self.appearance_embedding_dim)], dim=-1)
+
+        rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
+        return {FieldHeadNames.RGB: rgb}
+
+    # pylint: disable=arguments-differ
+    def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs":..., 1]) -> TensorType["bs":..., 1]:
+        """Returns only the density. Used primarily with the density grid.
+        Overwrite this function since density is time dependent now.
+
+        Args:
+            positions: the origin of the samples/frustums
+            times: the time of each position
+        """
+        ray_samples = RaySamples(
+            frustums=Frustums(
+                origins=positions,
+                directions=torch.ones_like(positions),
+                starts=torch.zeros_like(positions[..., :1]),
+                ends=torch.zeros_like(positions[..., :1]),
+                pixel_area=torch.ones_like(positions[..., :1]),
+            ),
+            times=times,
+        )
+        density, _ = self.get_density(ray_samples)
+        return density
+
+    def get_opacity(self, positions: TensorType["bs":..., 3], step_size, time_intervals=10) -> TensorType["bs":..., 1]:
+        """Returns the opacity for a position and time. Used primarily by the occupancy grid.
+        This will return the maximum opacity of the points in the space in a dynamic sequence.
+
+        Args:
+            positions: the positions to evaluate the opacity at.
+            step_size: the step size to use for the opacity evaluation.
+            time_intervals: sample density on N time stamps
+        """
+        # TODO: Converting opacity by time intervals is slow, and may lead to temporal artifacts.
+        #       (Maybe random sample time and EMA?)
+        opacity = []
+        for t in range(0, time_intervals):
+            density = self.density_fn(positions, t / (time_intervals - 1) * torch.ones_like(positions)[..., [0]])
+            opacity.append(density * step_size)
+        opacity = torch.stack(opacity, dim=0).max(dim=0).values
+        return opacity