nerfplayer with ngp and nerfacto (#1277)

* nerfplayer with ngp and nerfacto

* update comments

* update cpp extension build

* avoid Windows repeated CUDA compile

---------

Co-authored-by: machenmusik <machenmusik@users.noreply.github.com>

nerfstudio/configs/method_configs.py

@@ -33,6 +33,7 @@ from nerfstudio.data.datamanagers.variable_res_datamanager import (
 )
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
+from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig
 from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
 from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
     InstantNGPDataParserConfig,
@@ -49,6 +50,8 @@ from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
 from nerfstudio.models.mipnerf import MipNerfModel
 from nerfstudio.models.nerfacto import NerfactoModelConfig
+from nerfstudio.models.nerfplayer_nerfacto import NerfplayerNerfactoModelConfig
+from nerfstudio.models.nerfplayer_ngp import NerfplayerNGPModelConfig
 from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
 from nerfstudio.models.tensorf import TensoRFModelConfig
 from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig
@@ -67,6 +70,8 @@ descriptions = {
     "tensorf": "tensorf",
     "dnerf": "Dynamic-NeRF model. (slow)",
     "phototourism": "Uses the Phototourism data.",
+    "nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.",
+    "nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.",
 }
 
 method_configs["nerfacto"] = TrainerConfig(
@@ -325,6 +330,63 @@ method_configs["phototourism"] = TrainerConfig(
     vis="viewer",
 )
 
+method_configs["nerfplayer-nerfacto"] = TrainerConfig(
+    method_name="nerfplayer-nerfacto",
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
+    pipeline=VanillaPipelineConfig(
+        datamanager=DepthDataManagerConfig(
+            dataparser=DycheckDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NerfplayerNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+method_configs["nerfplayer-ngp"] = TrainerConfig(
+    method_name="nerfplayer-ngp",
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
+    pipeline=DynamicBatchPipelineConfig(
+        datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192),
+        model=NerfplayerNGPModelConfig(
+            eval_num_rays_per_chunk=8192,
+            contraction_type=ContractionType.AABB,
+            render_step_size=0.001,
+            max_num_samples_per_ray=48,
+            near_plane=0.01,
+        ),
+    ),
+    optimizers={
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        }
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=64000),
+    vis="viewer",
+)
+
 AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[  # Don't show unparseable (fixed) arguments in helptext.
     tyro.conf.FlagConversionOff[
         tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)

nerfstudio/field_components/cuda/_backend.py

@@ -27,9 +27,10 @@ BUILD_DIR = _get_build_directory(NAME, verbose=False)
 
 _C = None
 if cuda_toolkit_available():
-    if os.path.exists(os.path.join(BUILD_DIR, f"{NAME}.so")):
+    if os.listdir(BUILD_DIR) != []:
         # If the build exists, we assume the extension has been built
         # and we can load it.
+        print("nerfstudio field components: CUDA set up, loading (should be quick)")
         _C = load(
             name=NAME,
             sources=glob.glob(os.path.join(PATH, "csrc/*.cu")),

nerfstudio/field_components/temporal_grid.py

@@ -182,7 +182,7 @@ class TemporalGridEncoder(nn.Module):
         input_dim: int = 3,
         num_levels: int = 16,
         level_dim: int = 2,
-        per_level_scale: int = 2,
+        per_level_scale: float = 2.0,
         base_resolution: int = 16,
         log2_hashmap_size: int = 19,
         desired_resolution: Optional[int] = None,

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

nerfstudio/model_components/ray_samplers.py

@@ -503,6 +503,8 @@ class VolumetricSampler(Sampler):
             ),
             camera_indices=camera_indices,
         )
+        if ray_bundle.times is not None:
+            ray_samples.times = ray_bundle.times[ray_indices]
         return ray_samples, ray_indices
 
 

nerfstudio/models/nerfplayer_nerfacto.py

@@ -0,0 +1,251 @@
+# 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) implementation with nerfacto backbone.
+"""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass, field
+from typing import Dict, List, Type
+
+import numpy as np
+import torch
+from torchmetrics import PeakSignalNoiseRatio
+from torchmetrics.functional import structural_similarity_index_measure
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal
+
+from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.spatial_distortions import SceneContraction
+from nerfstudio.fields.nerfplayer_nerfacto_field import (
+    NerfplayerNerfactoField,
+    TemporalHashMLPDensityField,
+)
+from nerfstudio.model_components.losses import (
+    MSELoss,
+    interlevel_loss,
+    orientation_loss,
+    pred_normal_loss,
+)
+from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
+from nerfstudio.model_components.renderers import (
+    AccumulationRenderer,
+    DepthRenderer,
+    NormalsRenderer,
+    RGBRenderer,
+)
+from nerfstudio.model_components.scene_colliders import NearFarCollider
+from nerfstudio.models.base_model import Model
+from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
+
+
+@dataclass
+class NerfplayerNerfactoModelConfig(NerfactoModelConfig):
+    """Nerfplayer Model Config with Nerfacto backbone"""
+
+    _target: Type = field(default_factory=lambda: NerfplayerNerfactoModel)
+    near_plane: float = 0.05
+    """How far along the ray to start sampling."""
+    far_plane: float = 1000.0
+    """How far along the ray to stop sampling."""
+    background_color: Literal["random", "last_sample"] = "random"
+    """Whether to randomize the background color. (Random is reported to be better on DyCheck.)"""
+    num_levels: int = 16
+    """Hashing grid parameter."""
+    features_per_level: int = 2
+    """Hashing grid parameter."""
+    log2_hashmap_size: int = 18
+    """Hashing grid parameter."""
+    temporal_dim: int = 32
+    """Hashing grid parameter. A higher temporal dim means a higher temporal frequency."""
+    proposal_net_args_list: List[Dict] = field(
+        default_factory=lambda: [
+            {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 64},
+            {"hidden_dim": 16, "temporal_dim": 32, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256},
+        ]
+    )
+    """Arguments for the proposal density fields."""
+    distortion_loss_mult: float = 1e-2
+    """Distortion loss multiplier."""
+    temporal_tv_weight: float = 1
+    """Temporal TV balancing weight for feature channels."""
+    depth_weight: float = 1e-1
+    """depth loss balancing weight for feature channels."""
+
+
+class NerfplayerNerfactoModel(NerfactoModel):
+    """Nerfplayer model with Nerfacto backbone.
+
+    Args:
+        config: Nerfplayer configuration to instantiate model
+    """
+
+    config: NerfplayerNerfactoModelConfig
+
+    def populate_modules(self):
+        """Set the fields and modules."""
+        Model.populate_modules(self)
+
+        scene_contraction = SceneContraction(order=float("inf"))
+
+        # Fields
+        self.field = NerfplayerNerfactoField(
+            self.scene_box.aabb,
+            temporal_dim=self.config.temporal_dim,
+            num_levels=self.config.num_levels,
+            features_per_level=self.config.features_per_level,
+            log2_hashmap_size=self.config.log2_hashmap_size,
+            spatial_distortion=scene_contraction,
+            num_images=self.num_train_data,
+            use_pred_normals=self.config.predict_normals,
+            use_average_appearance_embedding=self.config.use_average_appearance_embedding,
+        )
+
+        self.density_fns = []
+        num_prop_nets = self.config.num_proposal_iterations
+        # Build the proposal network(s)
+        self.proposal_networks = torch.nn.ModuleList()
+        if self.config.use_same_proposal_network:
+            assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
+            prop_net_args = self.config.proposal_net_args_list[0]
+            network = TemporalHashMLPDensityField(
+                self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args
+            )
+            self.proposal_networks.append(network)
+            self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
+        else:
+            for i in range(num_prop_nets):
+                prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
+                network = TemporalHashMLPDensityField(
+                    self.scene_box.aabb,
+                    spatial_distortion=scene_contraction,
+                    **prop_net_args,
+                )
+                self.proposal_networks.append(network)
+            self.density_fns.extend([network.density_fn for network in self.proposal_networks])
+
+        # Samplers
+        update_schedule = lambda step: np.clip(
+            np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]),
+            1,
+            self.config.proposal_update_every,
+        )
+        self.proposal_sampler = ProposalNetworkSampler(
+            num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
+            num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
+            num_proposal_network_iterations=self.config.num_proposal_iterations,
+            single_jitter=self.config.use_single_jitter,
+            update_sched=update_schedule,
+        )
+
+        # Collider
+        self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
+
+        # renderers
+        self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
+        self.renderer_accumulation = AccumulationRenderer()
+        self.renderer_depth = DepthRenderer(method="expected")  # for depth loss
+        self.renderer_normals = NormalsRenderer()
+
+        # losses
+        self.rgb_loss = MSELoss()
+
+        # metrics
+        self.psnr = PeakSignalNoiseRatio(data_range=1.0)
+        self.ssim = structural_similarity_index_measure
+        self.lpips = LearnedPerceptualImagePatchSimilarity()
+        self.temporal_distortion = True  # for viewer
+
+    def get_outputs(self, ray_bundle: RayBundle):
+        assert ray_bundle.times is not None, "Time not provided."
+        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(
+            ray_bundle, density_fns=[functools.partial(f, times=ray_bundle.times) for f in self.density_fns]
+        )
+        field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
+        weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
+        weights_list.append(weights)
+        ray_samples_list.append(ray_samples)
+
+        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
+        accumulation = self.renderer_accumulation(weights=weights)
+
+        outputs = {
+            "rgb": rgb,
+            "accumulation": accumulation,
+            "depth": depth,
+        }
+
+        if self.config.predict_normals:
+            outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
+            outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
+
+        # These use a lot of GPU memory, so we avoid storing them for eval.
+        if self.training:
+            outputs["weights_list"] = weights_list
+            outputs["ray_samples_list"] = ray_samples_list
+
+        if self.training and self.config.predict_normals:
+            outputs["rendered_orientation_loss"] = orientation_loss(
+                weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
+            )
+
+            outputs["rendered_pred_normal_loss"] = pred_normal_loss(
+                weights.detach(),
+                field_outputs[FieldHeadNames.NORMALS].detach(),
+                field_outputs[FieldHeadNames.PRED_NORMALS],
+            )
+
+        for i in range(self.config.num_proposal_iterations):
+            outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
+
+        return outputs
+
+    def get_loss_dict(self, outputs, batch, metrics_dict=None):
+        loss_dict = {}
+        image = batch["image"].to(self.device)
+        loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
+        if self.training:
+            loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
+                outputs["weights_list"], outputs["ray_samples_list"]
+            )
+            assert metrics_dict is not None and "distortion" in metrics_dict
+            loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"]
+            if self.config.predict_normals:
+                # orientation loss for computed normals
+                loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean(
+                    outputs["rendered_orientation_loss"]
+                )
+
+                # ground truth supervision for normals
+                loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
+                    outputs["rendered_pred_normal_loss"]
+                )
+            if "depth_image" in batch.keys() and self.config.depth_weight > 0:
+                mask = (batch["depth_image"] != 0).view([-1])
+                loss_dict["depth_loss"] = 0
+                l = lambda x: self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean()
+                loss_dict["depth_loss"] = l(outputs["depth"][mask])
+                for i in range(self.config.num_proposal_iterations):
+                    loss_dict["depth_loss"] += l(outputs[f"prop_depth_{i}"][mask])
+            if self.config.temporal_tv_weight > 0:
+                loss_dict["temporal_tv_loss"] = self.field.mlp_base.get_temporal_tv_loss()
+                for net in self.proposal_networks:
+                    loss_dict["temporal_tv_loss"] += net.encoding.get_temporal_tv_loss()
+                loss_dict["temporal_tv_loss"] *= self.config.temporal_tv_weight
+        return loss_dict

nerfstudio/models/nerfplayer_ngp.py

@@ -0,0 +1,241 @@
+# 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.
+
+"""
+Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP backbone.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Type
+
+import nerfacc
+import torch
+from nerfacc import ContractionType
+from torch.nn import Parameter
+from torchmetrics import PeakSignalNoiseRatio
+from torchmetrics.functional import structural_similarity_index_measure
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal
+
+from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.fields.nerfplayer_ngp_field import NerfplayerNGPField
+from nerfstudio.model_components.losses import MSELoss
+from nerfstudio.model_components.ray_samplers import VolumetricSampler
+from nerfstudio.model_components.renderers import (
+    AccumulationRenderer,
+    DepthRenderer,
+    RGBRenderer,
+)
+from nerfstudio.models.base_model import Model
+from nerfstudio.models.instant_ngp import InstantNGPModelConfig, NGPModel
+from nerfstudio.utils import colors
+
+
+@dataclass
+class NerfplayerNGPModelConfig(InstantNGPModelConfig):
+    """NeRFPlayer Model Config with InstantNGP backbone.
+    Tips for tuning the performance:
+    1. If the scene is flickering, this is caused by unwanted high-freq on the temporal dimension.
+        Try reducing `temporal_dim` first, but don't be too small, otherwise the dynamic object is blurred.
+        Then try increasing the `temporal_tv_weight`. This is the loss for promoting smoothness among the
+        temporal channels.
+    2. If a faster rendering is preferred, then try reducing `log2_hashmap_size`. If more details are
+        wanted, try increasing `log2_hashmap_size`.
+    3. If the input cameras are of limited numbers, try reducing `num_levels`. `num_levels` is for
+        multi-resolution volume sampling, and has a similar behavior to the freq in NeRF. With a small
+        `num_levels`, a blurred rendering will be generated, but it is unlikely to overfit the training views.
+    """
+
+    _target: Type = field(default_factory=lambda: NerfplayerNGPModel)
+    temporal_dim: int = 64
+    """Hashing grid parameter. A higher temporal dim means a higher temporal frequency."""
+    num_levels: int = 16
+    """Hashing grid parameter."""
+    features_per_level: int = 2
+    """Hashing grid parameter."""
+    log2_hashmap_size: int = 17
+    """Hashing grid parameter."""
+    base_resolution: int = 16
+    """Hashing grid parameter."""
+    temporal_tv_weight: float = 1
+    """Temporal TV loss balancing weight for feature channels."""
+    depth_weight: float = 1e-1
+    """depth loss balancing weight for feature channels."""
+    train_background_color: Literal["random", "black", "white"] = "random"
+    """The training background color that is given to untrained areas."""
+    eval_background_color: Literal["random", "black", "white"] = "white"
+    """The training background color that is given to untrained areas."""
+    disable_viewing_dependent: bool = True
+    """Disable viewing dependent effects."""
+
+
+class NerfplayerNGPModel(NGPModel):
+    """NeRFPlayer Model with InstantNGP backbone.
+
+    Args:
+        config: NeRFPlayer NGP configuration to instantiate model
+    """
+
+    config: NerfplayerNGPModelConfig
+    field: NerfplayerNGPField
+
+    def populate_modules(self):
+        """Set the fields and modules."""
+        Model.populate_modules(self)
+
+        self.field = NerfplayerNGPField(
+            aabb=self.scene_box.aabb,
+            contraction_type=self.config.contraction_type,
+            use_appearance_embedding=self.config.use_appearance_embedding,
+            num_images=self.num_train_data,
+            temporal_dim=self.config.temporal_dim,
+            num_levels=self.config.num_levels,
+            features_per_level=self.config.features_per_level,
+            log2_hashmap_size=self.config.log2_hashmap_size,
+            base_resolution=self.config.base_resolution,
+            disable_viewing_dependent=self.config.disable_viewing_dependent,
+        )
+
+        self.scene_aabb = Parameter(self.scene_box.aabb.flatten(), requires_grad=False)
+
+        # Occupancy Grid
+        self.occupancy_grid = nerfacc.OccupancyGrid(
+            roi_aabb=self.scene_aabb,
+            resolution=self.config.grid_resolution,
+            contraction_type=self.config.contraction_type,
+        )
+
+        # Sampler
+        vol_sampler_aabb = self.scene_box.aabb if self.config.contraction_type == ContractionType.AABB else None
+        self.sampler = VolumetricSampler(
+            scene_aabb=vol_sampler_aabb,
+            occupancy_grid=self.occupancy_grid,
+            density_fn=self.field.density_fn,
+        )  # need to update the density_fn later during forward (for input time)
+
+        # renderers
+        self.train_background_color = self.config.train_background_color
+        self.eval_background_color = self.config.eval_background_color
+        if self.config.train_background_color in ["white", "black"]:
+            self.train_background_color = colors.COLORS_DICT[self.config.train_background_color]
+        if self.config.eval_background_color in ["white", "black"]:
+            self.eval_background_color = colors.COLORS_DICT[self.config.eval_background_color]
+
+        self.renderer_rgb = RGBRenderer()  # will update bgcolor later during forward
+        self.renderer_accumulation = AccumulationRenderer()
+        self.renderer_depth = DepthRenderer(method="expected")
+
+        # losses
+        self.rgb_loss = MSELoss()
+
+        # metrics
+        self.psnr = PeakSignalNoiseRatio(data_range=1.0)
+        self.ssim = structural_similarity_index_measure
+        self.lpips = LearnedPerceptualImagePatchSimilarity()
+        self.temporal_distortion = True  # for viewer
+
+    def get_outputs(self, ray_bundle: RayBundle):
+        num_rays = len(ray_bundle)
+
+        # update the density_fn of the sampler so that the density is time aware
+        self.sampler.density_fn = lambda x: self.field.density_fn(x, ray_bundle.times)
+        with torch.no_grad():
+            ray_samples, ray_indices = self.sampler(
+                ray_bundle=ray_bundle,
+                near_plane=self.config.near_plane,
+                far_plane=self.config.far_plane,
+                render_step_size=self.config.render_step_size,
+                cone_angle=self.config.cone_angle,
+            )
+
+        field_outputs = self.field(ray_samples)
+
+        # accumulation
+        packed_info = nerfacc.pack_info(ray_indices, num_rays)
+        weights = nerfacc.render_weight_from_density(
+            packed_info=packed_info,
+            sigmas=field_outputs[FieldHeadNames.DENSITY],
+            t_starts=ray_samples.frustums.starts,
+            t_ends=ray_samples.frustums.ends,
+        )
+
+        # update bgcolor in the renderer; usually random color for training and fixed color for inference
+        if self.training:
+            self.renderer_rgb.background_color = self.train_background_color
+        else:
+            self.renderer_rgb.background_color = self.eval_background_color
+        rgb = self.renderer_rgb(
+            rgb=field_outputs[FieldHeadNames.RGB],
+            weights=weights,
+            ray_indices=ray_indices,
+            num_rays=num_rays,
+        )
+        depth = self.renderer_depth(
+            weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays
+        )
+        accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)
+        alive_ray_mask = accumulation.squeeze(-1) > 0
+
+        outputs = {
+            "rgb": rgb,
+            "accumulation": accumulation,
+            "depth": depth,
+            "alive_ray_mask": alive_ray_mask,  # the rays we kept from sampler
+            "num_samples_per_ray": packed_info[:, 1],
+        }
+        # adding them to outputs for calculating losses
+        if self.training and self.config.depth_weight > 0:
+            outputs["ray_indices"] = ray_indices
+            outputs["ray_samples"] = ray_samples
+            outputs["weights"] = weights
+            outputs["sigmas"] = field_outputs[FieldHeadNames.DENSITY]
+        return outputs
+
+    def get_loss_dict(self, outputs, batch, metrics_dict=None):
+        image = batch["image"].to(self.device)
+        mask = outputs["alive_ray_mask"]
+        rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask])
+        loss_dict = {"rgb_loss": rgb_loss}
+        if "depth_image" in batch.keys() and self.config.depth_weight > 0:
+            mask = batch["depth_image"] != 0
+            # First we calculate the depth value, just like most of the papers.
+            loss_dict["depth_loss"] = (outputs["depth"][mask] - batch["depth_image"][mask]).abs().mean()
+            # But this is not enough -- it will lead to fog like reconstructions, even with depth supervision.
+            # Because this loss only cares about the mean value.
+            # (Feel free to try it and see the fog-like effects on DyCheck by commenting out the following loss.)
+            # > The fog can be effectively penalized by multiview inputs (i.e., from another viewing point).
+            # > However, it is hard to get multiple views under the setting of dynamic scenes.
+            # > A new view always indicates another camera, which is expensive and brings sync problems.
+            # In DyCheck (https://arxiv.org/abs/2210.13445), surface sparsity regularizer
+            # (i.e., `distortion_loss` in nerfstudio) is used to make surface tight.
+            # The `distortion_loss` can be used here as well, but personally find it hard to implement...
+            # (Due to volume sampling, seems that cuda kernels are needed for efficiently computing the loss.)
+            # (Try nerfplayer with nerfacto backbone for distortion loss.)
+            # Instead, directly penalizing the empty space is found effective here. Perhaps due to a more
+            # "clear" loss (as it is directly applied to the network outputs, rather than post-processed values).
+            # But such a loss also has drawbacks: it tends to overfit wrong (or noise) presented in the depth map.
+            # Some structures are floating in the air with this loss...
+            gt_depth_packed = batch["depth_image"][outputs["ray_indices"]]
+            steps = (outputs["ray_samples"].frustums.starts + outputs["ray_samples"].frustums.ends) / 2
+            # empty area should not be too close to the given depth, so lets add a margin to the gt depth
+            margin = (self.scene_box.aabb.max() - self.scene_box.aabb.min()) / 128
+            density_min_mask = (gt_depth_packed - steps > margin) & (gt_depth_packed != 0)
+            loss_dict["depth_loss"] += (outputs["sigmas"][density_min_mask[..., 0]].pow(2)).mean() * 1e-2
+            loss_dict["depth_loss"] *= self.config.depth_weight
+        if self.config.temporal_tv_weight > 0:
+            loss_dict["temporal_tv_loss"] = self.config.temporal_tv_weight * self.field.mlp_base.get_temporal_tv_loss()
+        return loss_dict

nerfstudio/viewer/server/viewer_utils.py

@@ -196,11 +196,16 @@ class CheckThread(threading.Thread):
         while not self.state.check_done_render:
             # check camera
             data = self.state.vis["renderingState/camera"].read()
+            render_time = self.state.vis["renderingState/render_time"].read()
             if data is not None:
                 camera_object = data["object"]
-                if self.state.prev_camera_matrix is None or (
-                    not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix)
-                    and not self.state.prev_moving
+                if (
+                    self.state.prev_camera_matrix is None
+                    or (
+                        not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix)
+                        and not self.state.prev_moving
+                    )
+                    or (render_time is not None and render_time != self.state.prev_render_time)
                 ):
                     self.state.check_interrupt_vis = True
                     self.state.prev_moving = True
@@ -280,6 +285,7 @@ class ViewerState:
 
         # viewer specific variables
         self.prev_camera_matrix = None
+        self.prev_render_time = 0
         self.prev_output_type = OutputTypes.INIT
         self.prev_colormap_type = ColormapTypes.INIT
         self.prev_moving = False
@@ -553,12 +559,23 @@ class ViewerState:
             return None
 
         camera_object = data["object"]
+        render_time = self.vis["renderingState/render_time"].read()
 
-        if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix):
-            self.camera_moving = False
+        if render_time is not None:
+            if (
+                self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix)
+            ) and (self.prev_render_time == render_time):
+                self.camera_moving = False
+            else:
+                self.prev_camera_matrix = camera_object["matrix"]
+                self.prev_render_time = render_time
+                self.camera_moving = True
         else:
-            self.prev_camera_matrix = camera_object["matrix"]
-            self.camera_moving = True
+            if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix):
+                self.camera_moving = False
+            else:
+                self.prev_camera_matrix = camera_object["matrix"]
+                self.camera_moving = True
 
         output_type = self.vis["renderingState/output_choice"].read()
         if output_type is None:
@@ -901,10 +918,6 @@ class ViewerState:
             dim=0,
         )
 
-        times = self.vis["renderingState/render_time"].read()
-        if times is not None:
-            times = torch.tensor([float(times)])
-
         camera_type_msg = camera_object["camera_type"]
         if camera_type_msg == "perspective":
             camera_type = CameraType.PERSPECTIVE
@@ -922,11 +935,11 @@ class ViewerState:
             cy=intrinsics_matrix[1, 2],
             camera_type=camera_type,
             camera_to_worlds=camera_to_world[None, ...],
-            times=times,
+            times=torch.tensor([float(self.prev_render_time)]),
         )
         camera = camera.to(graph.device)
 
-        camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=graph.render_aabb)
+        camera_ray_bundle = camera.generate_rays(camera_indices=0)
 
         graph.eval()
 

tests/test_train.py

@@ -23,6 +23,8 @@ BLACKLIST = [
     "nerfacto",
     "phototourism",
     "depth-nerfacto",
+    "nerfplayer-ngp",
+    "nerfplayer-nerfacto",
 ]