one commit

nerfstudio/cameras/cameras.py

@@ -766,7 +766,7 @@ class Cameras(TensorDataclass):
         Returns:
             Pinhole camera intrinsics matrices
         """
-        K = torch.zeros((*self.shape, 3, 3), dtype=torch.float32)
+        K = torch.zeros((*self.shape, 3, 3), dtype=torch.float32, device=self.device)
         K[..., 0, 0] = self.fx.squeeze(-1)
         K[..., 1, 1] = self.fy.squeeze(-1)
         K[..., 0, 2] = self.cx.squeeze(-1)

nerfstudio/configs/experiment_config.py

@@ -50,6 +50,8 @@ class ExperimentConfig(InstantiateConfig):
     """Method name. Required to set in python or via cli"""
     experiment_name: Optional[str] = None
     """Experiment name. If None, will automatically be set to dataset name"""
+    project_name: Optional[str] = "nerfstudio-project"
+    """Project name."""
     timestamp: str = "{timestamp}"
     """Experiment timestamp."""
     machine: MachineConfig = MachineConfig()

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -16,9 +16,10 @@
 from __future__ import annotations
 
 import math
+import os
 from dataclasses import dataclass, field
 from pathlib import Path, PurePath
-from typing import Optional, Type
+from typing import Optional, Tuple, Type
 
 import numpy as np
 import torch
@@ -60,8 +61,14 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """The method to use to center the poses."""
     auto_scale_poses: bool = True
     """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
+    eval_mode: Literal["train-split-fraction", "eval-frame-index"] = "train-split-fraction"
+    """The method to use for splitting the dataset into train and eval."""
     train_split_fraction: float = 0.9
     """The fraction of images to use for training. The remaining images are for eval."""
+    train_frame_indices: Tuple[int, ...] = (0,)
+    """The index of the frames to use for train."""
+    eval_frame_indices: Tuple[int, ...] = (1,)
+    """The index of the frames to use for eval."""
     depth_unit_scale_factor: float = 1e-3
     """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
 
@@ -110,7 +117,17 @@ class Nerfstudio(DataParser):
         width = []
         distort = []
 
+        # sort the frames by fname
+        # they do this in mipnerf360 code
+        fnames = []
         for frame in meta["frames"]:
+            filepath = PurePath(frame["file_path"])
+            fname = self._get_fname(filepath, data_dir)
+            fnames.append(fname)
+        inds = np.argsort(fnames)
+        frames = [meta["frames"][ind] for ind in inds]
+
+        for frame in frames:
             filepath = PurePath(frame["file_path"])
             fname = self._get_fname(filepath, data_dir)
             if not fname.exists():
@@ -184,30 +201,61 @@ class Nerfstudio(DataParser):
         You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json.
         """
 
-        has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test"))
-        if f"{split}_filenames" in meta:
-            # Validate split first
-            split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"])
-            unmatched_filenames = split_filenames.difference(image_filenames)
-            if unmatched_filenames:
-                raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.")
+        if self.config.eval_mode == "train-split-fraction":
+            has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test"))
+            if f"{split}_filenames" in meta:
+                # Validate split first
+                split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"])
+                unmatched_filenames = split_filenames.difference(image_filenames)
+                if unmatched_filenames:
+                    raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.")
 
-            indices = [i for i, path in enumerate(image_filenames) if path in split_filenames]
-            CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}")
-            indices = np.array(indices, dtype=np.int32)
-        elif has_split_files_spec:
-            raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.")
-        else:
-            # filter image_filenames and poses based on train/eval split percentage
+                indices = [i for i, path in enumerate(image_filenames) if path in split_filenames]
+                CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}")
+                indices = np.array(indices, dtype=np.int32)
+            elif has_split_files_spec:
+                raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.")
+            else:
+                # filter image_filenames and poses based on train/eval split percentage
+                num_images = len(image_filenames)
+                num_train_images = math.ceil(num_images * self.config.train_split_fraction)
+                num_eval_images = num_images - num_train_images
+                i_all = np.arange(num_images)
+                i_train = np.linspace(
+                    0, num_images - 1, num_train_images, dtype=int
+                )  # equally spaced training images starting and ending at 0 and num_images-1
+                i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
+                assert len(i_eval) == num_eval_images
+                if split == "train":
+                    indices = i_train
+                elif split in ["val", "test"]:
+                    indices = i_eval
+                else:
+                    raise ValueError(f"Unknown dataparser split {split}")
+        elif self.config.eval_mode == "eval-frame-index":
+            # keep around some metadata
+            eval_frame_index_0_metadata = []
+            # filter image_filenames and poses based on train and eval frame indices
             num_images = len(image_filenames)
-            num_train_images = math.ceil(num_images * self.config.train_split_fraction)
-            num_eval_images = num_images - num_train_images
+            basenames = [os.path.basename(image_filename) for image_filename in image_filenames]
             i_all = np.arange(num_images)
-            i_train = np.linspace(
-                0, num_images - 1, num_train_images, dtype=int
-            )  # equally spaced training images starting and ending at 0 and num_images-1
-            i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
-            assert len(i_eval) == num_eval_images
+            i_train = []
+            i_eval = []
+            for idx, basename in zip(i_all, basenames):
+                # check the frame index
+                if len(basename.split("_")) == 2:
+                    frame_index = 0
+                else:
+                    frame_index = int(basename.split("_")[1])
+                if frame_index in self.config.train_frame_indices:
+                    i_train.append(idx)
+                    if frame_index == 0:
+                        # set 1 where frame_index is 0
+                        eval_frame_index_0_metadata.append(1)
+                    else:
+                        eval_frame_index_0_metadata.append(0)
+                if frame_index in self.config.eval_frame_indices:
+                    i_eval.append(idx)
             if split == "train":
                 indices = i_train
             elif split in ["val", "test"]:
@@ -290,6 +338,14 @@ class Nerfstudio(DataParser):
         assert self.downscale_factor is not None
         cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)
 
+        metadata = {
+            "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
+            "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
+        }
+
+        if self.config.eval_mode == "eval-frame-index":
+            metadata["eval_frame_index_0_metadata"] = torch.tensor(eval_frame_index_0_metadata, dtype=torch.long)
+
         if "applied_transform" in meta:
             applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
             transform_matrix = transform_matrix @ torch.cat(
@@ -306,10 +362,7 @@ class Nerfstudio(DataParser):
             mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
             dataparser_scale=scale_factor,
             dataparser_transform=transform_matrix,
-            metadata={
-                "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
-                "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
-            },
+            metadata=metadata,
         )
         return dataparser_outputs
 

nerfstudio/data/utils/dataloaders.py

@@ -211,7 +211,6 @@ class FixedIndicesEvalDataloader(EvalDataloader):
         self.count = 0
 
     def __iter__(self):
-        self.count = 0
         return self
 
     def __next__(self):
@@ -238,17 +237,12 @@ class RandIndicesEvalDataloader(EvalDataloader):
         **kwargs,
     ):
         super().__init__(input_dataset, device, **kwargs)
-        self.count = 0
 
     def __iter__(self):
-        self.count = 0
         return self
 
     def __next__(self):
-        if self.count < 1:
-            image_indices = range(self.cameras.size)
-            image_idx = random.choice(image_indices)
-            ray_bundle, batch = self.get_data_from_image_idx(image_idx)
-            self.count += 1
-            return ray_bundle, batch
-        raise StopIteration
+        # choose a random image index
+        image_idx = random.randint(0, len(self.cameras) - 1)
+        ray_bundle, batch = self.get_data_from_image_idx(image_idx)
+        return ray_bundle, batch

nerfstudio/engine/callbacks.py

@@ -37,6 +37,8 @@ class TrainingCallbackAttributes:
     """gradient scalers"""
     pipeline: Optional[InitVar]
     """reference to training pipeline"""
+    viewer_state: Optional[InitVar]
+    """reference to viewer state"""
 
 
 class TrainingCallbackLocation(Enum):

nerfstudio/engine/trainer.py

@@ -126,7 +126,6 @@ class Trainer:
         # directory to save checkpoints
         self.checkpoint_dir: Path = config.get_checkpoint_dir()
         CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}")
-
         self.viewer_state = None
 
     def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None:
@@ -162,13 +161,18 @@ class Trainer:
                 optimizers=self.optimizers,  # type: ignore
                 grad_scaler=self.grad_scaler,  # type: ignore
                 pipeline=self.pipeline,  # type: ignore
+                viewer_state=self.viewer_state,  # type: ignore
             )
         )
 
         # set up writers/profilers if enabled
         writer_log_path = self.base_dir / self.config.logging.relative_log_dir
         writer.setup_event_writer(
-            self.config.is_wandb_enabled(), self.config.is_tensorboard_enabled(), log_dir=writer_log_path
+            self.config.is_wandb_enabled(),
+            self.config.is_tensorboard_enabled(),
+            log_dir=writer_log_path,
+            experiment_name=self.config.experiment_name,
+            project_name=self.config.project_name,
         )
         writer.setup_local_writer(
             self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages
@@ -288,6 +292,7 @@ class Trainer:
         self.viewer_state.init_scene(
             dataset=self.pipeline.datamanager.train_dataset,
             start_train=self.config.viewer.start_train,
+            eval_dataset=self.pipeline.datamanager.eval_dataset,
         )
         if not self.config.viewer.start_train:
             self._always_render(self._start_step)
@@ -335,7 +340,7 @@ class Trainer:
         if load_dir is not None:
             load_step = self.config.load_step
             if load_step is None:
-                print("Loading latest checkpoint from load_dir")
+                print("Loading latest Nerfstudio checkpoint from load_dir...")
                 # NOTE: this is specific to the checkpoint name format
                 load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1]
             load_path: Path = load_dir / f"step-{load_step:09d}.ckpt"
@@ -346,9 +351,9 @@ class Trainer:
             self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])
             self.optimizers.load_optimizers(loaded_state["optimizers"])
             self.grad_scaler.load_state_dict(loaded_state["scalers"])
-            CONSOLE.print(f"done loading checkpoint from {load_path}")
+            CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}")
         else:
-            CONSOLE.print("No checkpoints to load, training from scratch")
+            CONSOLE.print("No Nerfstudio checkpoint to load, so training from scratch.")
 
     @check_main_thread
     def save_checkpoint(self, step: int) -> None:

nerfstudio/fields/visibility_field.py

@@ -0,0 +1,90 @@
+# Copyright 2022 The Plenoptix 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.
+
+"""Visibility Field"""
+
+import torch
+from torch import nn
+
+from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.cameras.rays import RaySamples
+
+
+class VisibilityField(nn.Module):
+    """Visibility Field"""
+
+    def __init__(self, cameras: Cameras) -> None:
+        super().__init__()
+        # training camera tranforms
+        # TODO: use optimized cameras
+        self.c2ws = cameras.camera_to_worlds
+        self.c2whs = torch.cat([self.c2ws, torch.zeros_like(self.c2ws[:, :1, :])], dim=1)
+        self.c2whs[:, 3, 3] = 1.0
+        self.w2chs = torch.inverse(self.c2whs)
+        self.K = cameras.get_intrinsics_matrices()
+        self.image_height = cameras.height
+        self.image_width = cameras.width
+
+    @torch.no_grad()
+    def forward(self, ray_samples: RaySamples, camera_chunk_size=50, ray_chunk_size=4096):
+        """
+        Args:
+            ray_samples: Ray samples.
+            camera_chunk_size: Number of cameras to process at once to avoid memory issues.
+            ray_chunk_size: Number of rays to process at once to avoid memory issues.
+        Returns:
+        """
+        # get positions
+        positions = ray_samples.frustums.get_positions()  # [N, S, 3]
+        # project positions into each camera
+        # move to homogeneous coordinates
+        positions = torch.cat([positions, torch.ones_like(positions[..., :1])], dim=-1)
+        N, S, _ = positions.shape
+        B = self.w2chs.shape[0]  # num cameras
+        p = positions.view(N * S, 4).transpose(0, 1).unsqueeze(0)  # [1, 4, N*S]
+        p = p.expand(B, *p.shape[1:])  # [B, 4, N*S]
+
+        num_views = torch.zeros([N, S, 1], device=positions.device)
+        for i in range(0, B, camera_chunk_size):
+            ccs = min(camera_chunk_size, B - i)
+            for j in range(0, N, ray_chunk_size):
+                rcs = min(ray_chunk_size, N - j)
+
+                ptemp = p.reshape(B, 4, N, S)[i : i + ccs, :, j : j + rcs, :].reshape(ccs, 4, rcs * S)
+                cam_coords = torch.bmm(self.w2chs[i : i + ccs, :], ptemp)
+
+                # flip y and z axes
+                cam_coords[:, 1, :] *= -1
+                cam_coords[:, 2, :] *= -1
+
+                z = cam_coords[:, 2:3, :].transpose(1, 2).view(ccs, rcs, S, 1)  # [CS, RCS, S, 1]
+                mask_z = z[..., 0] > 0
+
+                # divide by z
+                cam_coords = cam_coords[:, :3, :] / cam_coords[:, 2:3, :]
+
+                cam_points = torch.bmm(self.K[i : i + ccs], cam_coords)
+
+                pixel_coords = cam_points[:, :2, :].transpose(1, 2).view(ccs, rcs, S, 2)  # [CS, RCS, S, 2]
+                x = pixel_coords[..., 0]
+                y = pixel_coords[..., 1]
+                mask_x = (x >= 0) & (x < self.image_width.view(B, 1, 1)[i : i + ccs])
+                mask_y = (y >= 0) & (y < self.image_height.view(B, 1, 1)[i : i + ccs])
+                mask = mask_x & mask_y & mask_z
+                # sum over the batch dimension
+                nv = mask.sum(dim=0).unsqueeze(-1)
+                # nv is [N, S, 1] # this is the number of camera frustums that the point belongs to
+                # for this particular chunk of cameras
+                num_views[j : j + rcs] += nv
+        return num_views

nerfstudio/model_components/ray_samplers.py

@@ -560,10 +560,11 @@ class ProposalNetworkSampler(Sampler):
         self,
         ray_bundle: Optional[RayBundle] = None,
         density_fns: Optional[List[Callable]] = None,
-    ) -> Tuple[RaySamples, List, List]:
+    ) -> Tuple[RaySamples, List, List, List]:
         assert ray_bundle is not None
         assert density_fns is not None
 
+        densities_list = []
         weights_list = []
         ray_samples_list = []
 
@@ -591,13 +592,14 @@ class ProposalNetworkSampler(Sampler):
                     with torch.no_grad():
                         density = density_fns[i_level](ray_samples.frustums.get_positions())
                 weights = ray_samples.get_weights(density)
+                densities_list.append(density)
                 weights_list.append(weights)  # (num_rays, num_samples)
                 ray_samples_list.append(ray_samples)
         if updated:
             self._steps_since_update = 0
 
         assert ray_samples is not None
-        return ray_samples, weights_list, ray_samples_list
+        return ray_samples, densities_list, weights_list, ray_samples_list
 
 
 class NeuSSampler(Sampler):

nerfstudio/model_components/renderers.py

@@ -244,6 +244,7 @@ class DepthRenderer(nn.Module):
         ray_samples: RaySamples,
         ray_indices: Optional[TensorType["num_samples"]] = None,
         num_rays: Optional[int] = None,
+        quantity: Optional[TensorType[..., "num_samples", 1]] = None,
     ) -> TensorType[..., 1]:
         """Composite samples along ray and calculate depths.
 
@@ -252,35 +253,37 @@ class DepthRenderer(nn.Module):
             ray_samples: Set of ray samples.
             ray_indices: Ray index for each sample, used when samples are packed.
             num_rays: Number of rays, used when samples are packed.
+            quantity: If specified, then take the median of this quantity instead of the depth.
 
         Returns:
             Outputs of depth values.
         """
 
-        if self.method == "median":
-            steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
+        if quantity is None:
+            # set the quantity to steps
+            quantity = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
 
+        if self.method == "median":
             if ray_indices is not None and num_rays is not None:
                 raise NotImplementedError("Median depth calculation is not implemented for packed samples.")
             cumulative_weights = torch.cumsum(weights[..., 0], dim=-1)  # [..., num_samples]
             split = torch.ones((*weights.shape[:-2], 1), device=weights.device) * 0.5  # [..., 1]
             median_index = torch.searchsorted(cumulative_weights, split, side="left")  # [..., 1]
-            median_index = torch.clamp(median_index, 0, steps.shape[-2] - 1)  # [..., 1]
-            median_depth = torch.gather(steps[..., 0], dim=-1, index=median_index)  # [..., 1]
+            median_index = torch.clamp(median_index, 0, quantity.shape[-2] - 1)  # [..., 1]
+            median_depth = torch.gather(quantity[..., 0], dim=-1, index=median_index)  # [..., 1]
             return median_depth
         if self.method == "expected":
             eps = 1e-10
-            steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
 
             if ray_indices is not None and num_rays is not None:
                 # Necessary for packed samples from volumetric ray sampler
-                depth = nerfacc.accumulate_along_rays(weights, ray_indices, steps, num_rays)
+                depth = nerfacc.accumulate_along_rays(weights, ray_indices, quantity, num_rays)
                 accumulation = nerfacc.accumulate_along_rays(weights, ray_indices, None, num_rays)
                 depth = depth / (accumulation + eps)
             else:
-                depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps)
+                depth = torch.sum(weights * quantity, dim=-2) / (torch.sum(weights, -2) + eps)
 
-            depth = torch.clip(depth, steps.min(), steps.max())
+            depth = torch.clip(depth, quantity.min(), quantity.max())
 
             return depth
 

nerfstudio/models/nerfacto.py

@@ -19,7 +19,7 @@ NeRF implementation that combines many recent advancements.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type
 
 import numpy as np
 import torch
@@ -39,6 +39,7 @@ from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.visibility_field import VisibilityField
 from nerfstudio.model_components.losses import (
     MSELoss,
     distortion_loss,
@@ -126,6 +127,8 @@ class NerfactoModelConfig(ModelConfig):
     """Whether use single jitter or not for the proposal networks."""
     predict_normals: bool = False
     """Whether to predict normals or not."""
+    depth_method: Literal["median", "expected"] = "median"
+    """Median or expected depth from the depth renderer."""
     disable_scene_contraction: bool = False
     """Whether to disable scene contraction or not."""
 
@@ -163,6 +166,10 @@ class NerfactoModel(Model):
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,
         )
 
+        # this can be set by the pipeline
+        # is set, then we create a visibility outputs
+        self.visibility_field: Optional[VisibilityField] = None
+
         self.density_fns = []
         num_prop_nets = self.config.num_proposal_iterations
         # Build the proposal network(s)
@@ -211,8 +218,10 @@ class NerfactoModel(Model):
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
-        self.renderer_depth = DepthRenderer()
+        self.renderer_depth = DepthRenderer(method=self.config.depth_method)
         self.renderer_normals = NormalsRenderer()
+        self.renderer_depth_expected = DepthRenderer(method="expected")
+        self.renderer_depth_median = DepthRenderer(method="median")
 
         # shaders
         self.normals_shader = NormalsShader()
@@ -263,9 +272,12 @@ class NerfactoModel(Model):
         return callbacks
 
     def get_outputs(self, ray_bundle: RayBundle):
-        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
+        ray_samples, densities_list, weights_list, ray_samples_list = self.proposal_sampler(
+            ray_bundle, density_fns=self.density_fns
+        )
         field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
         weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
+        densities_list.append(field_outputs[FieldHeadNames.DENSITY])
         weights_list.append(weights)
         ray_samples_list.append(ray_samples)
 
@@ -279,6 +291,22 @@ class NerfactoModel(Model):
             "depth": depth,
         }
 
+        if self.visibility_field is not None:
+            assert isinstance(self.visibility_field, VisibilityField), "self.visibility_field must be a VisibilityField"
+            visibility_samples = self.visibility_field.forward(ray_samples_list[-1])
+            visibility = self.renderer_depth_expected(
+                weights=weights, ray_samples=ray_samples, quantity=visibility_samples
+            )
+            outputs["visibility_count"] = visibility.long()
+            visibility = visibility.float() / len(self.visibility_field.c2ws)  # [0, 1] 1 is seen by all cameras
+            outputs["visibility"] = visibility
+            visibility_median = self.renderer_depth_median(
+                weights=weights, ray_samples=ray_samples, quantity=visibility_samples
+            )
+            outputs["visibility_median_count"] = visibility_median.long()
+            visibility_median = visibility_median.float() / len(self.visibility_field.c2ws)  # [0, 1]
+            outputs["visibility_median"] = visibility_median
+
         if self.config.predict_normals:
             normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
             pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
@@ -286,6 +314,7 @@ class NerfactoModel(Model):
             outputs["pred_normals"] = self.normals_shader(pred_normals)
         # These use a lot of GPU memory, so we avoid storing them for eval.
         if self.training:
+            outputs["densities_list"] = densities_list
             outputs["weights_list"] = weights_list
             outputs["ray_samples_list"] = ray_samples_list
 
@@ -364,6 +393,12 @@ class NerfactoModel(Model):
 
         images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
 
+        # add visibility
+        if "visibility" in outputs:
+            vis = colormaps.apply_colormap(outputs["visibility"])
+            combined_vis = torch.cat([vis], dim=1)
+            images_dict["visibility"] = combined_vis
+
         for i in range(self.config.num_proposal_iterations):
             key = f"prop_depth_{i}"
             prop_depth_i = colormaps.apply_depth_colormap(

nerfstudio/process_data/process_data_utils.py

@@ -180,23 +180,27 @@ def convert_video_to_images(
 def copy_images_list(
     image_paths: List[Path],
     image_dir: Path,
+    image_prefix: str = "frame_",
     crop_border_pixels: Optional[int] = None,
     crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
     verbose: bool = False,
+    keep_image_dir: bool = False,
 ) -> List[Path]:
     """Copy all images in a list of Paths. Useful for filtering from a directory.
     Args:
         image_paths: List of Paths of images to copy to a new directory.
         image_dir: Path to the output directory.
+        image_prefix: Prefix for the image filenames.
         crop_border_pixels: If not None, crops each edge by the specified number of pixels.
         crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right)
         verbose: If True, print extra logging.
+        keep_image_dir: If True, keep the original directory.
     Returns:
         A list of the copied image Paths.
     """
 
     # Remove original directory only if we provide a proper image folder path
-    if image_dir.is_dir() and len(image_paths):
+    if image_dir.is_dir() and len(image_paths) and not keep_image_dir:
         shutil.rmtree(image_dir, ignore_errors=True)
         image_dir.mkdir(exist_ok=True, parents=True)
 
@@ -206,13 +210,13 @@ def copy_images_list(
     for idx, image_path in enumerate(image_paths):
         if verbose:
             CONSOLE.log(f"Copying image {idx + 1} of {len(image_paths)}...")
-        copied_image_path = image_dir / f"frame_{idx + 1:05d}{image_path.suffix}"
+        copied_image_path = image_dir / f"{image_prefix}{idx + 1:05d}{image_path.suffix}"
         shutil.copy(image_path, copied_image_path)
         copied_image_paths.append(copied_image_path)
 
     if crop_border_pixels is not None:
         file_type = image_paths[0].suffix
-        filename = f"frame_%05d{file_type}"
+        filename = f"{image_prefix}%05d{file_type}"
         crop = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}"
         ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop} "{image_dir / filename}"'
         run_command(ffmpeg_cmd, verbose=verbose)
@@ -288,7 +292,12 @@ def copy_and_upscale_polycam_depth_maps_list(
 
 
 def copy_images(
-    data: Path, image_dir: Path, verbose, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
+    data: Path,
+    image_dir: Path,
+    verbose,
+    image_prefix: str = "frame_",
+    crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
+    keep_image_dir: bool = False,
 ) -> OrderedDict[Path, Path]:
     """Copy images from a directory to a new directory.
 
@@ -296,7 +305,9 @@ def copy_images(
         data: Path to the directory of images.
         image_dir: Path to the output directory.
         verbose: If True, print extra logging.
+        image_prefix: Prefix for the image filenames.
         crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right)
+        keep_image_dir: If True, keep the original directory.
     Returns:
         The mapping from the original filenames to the new ones.
     """
@@ -308,7 +319,12 @@ def copy_images(
             sys.exit(1)
 
         copied_images = copy_images_list(
-            image_paths=image_paths, image_dir=image_dir, crop_factor=crop_factor, verbose=verbose
+            image_paths=image_paths,
+            image_dir=image_dir,
+            image_prefix=image_prefix,
+            crop_factor=crop_factor,
+            verbose=verbose,
+            keep_image_dir=keep_image_dir,
         )
         return OrderedDict((original_path, new_path) for original_path, new_path in zip(image_paths, copied_images))
 

nerfstudio/utils/writer.py

@@ -188,7 +188,13 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages
 
 
 @check_main_thread
-def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log_dir: Path) -> None:
+def setup_event_writer(
+    is_wandb_enabled: bool,
+    is_tensorboard_enabled: bool,
+    log_dir: Path,
+    experiment_name: str,
+    project_name: str = "nerfstudio-project",
+) -> None:
     """Initialization of all event writers specified in config
 
     Args:
@@ -198,7 +204,7 @@ def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log
     """
     using_event_writer = False
     if is_wandb_enabled:
-        curr_writer = WandbWriter(log_dir=log_dir)
+        curr_writer = WandbWriter(log_dir=log_dir, experiment_name=experiment_name, project_name=project_name)
         EVENT_WRITERS.append(curr_writer)
         using_event_writer = True
     if is_tensorboard_enabled:
@@ -282,8 +288,8 @@ class TimeWriter:
 class WandbWriter(Writer):
     """WandDB Writer Class"""
 
-    def __init__(self, log_dir: Path):
-        wandb.init(project="nerfstudio-project", dir=str(log_dir), reinit=True, name=log_dir.name)
+    def __init__(self, log_dir: Path, experiment_name: str, project_name: str = "nerfstudio-project"):
+        wandb.init(project=project_name, dir=str(log_dir), reinit=True, name=experiment_name)
 
     def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
         image = torch.permute(image, (2, 0, 1))

nerfstudio/viewer/app/src/modules/Scene/Scene.jsx

@@ -19,6 +19,8 @@ const msgpack = require('msgpack-lite');
 
 const SCENE_BOX_NAME = 'Scene Box';
 const CAMERAS_NAME = 'Training Cameras';
+const EVAL_CAMERAS_NAME = 'Evaluation Cameras';
+const CUBES_NAME = 'Cubes';
 
 export function get_scene_tree() {
   const scene = new THREE.Scene();
@@ -286,6 +288,80 @@ export function get_scene_tree() {
   };
   subscribe_to_changes(selector_fn_cameras, fn_value_cameras);
 
+  // draw eval camera
+  // NOTE: this has some issues right now! it won't
+  // update the camera on an individual change w/o deleting first
+  const selector_fn_eval_cameras = (state) => {
+    return state.sceneState.eval_cameras;
+  };
+  const fn_value_eval_cameras = (previous, current) => {
+    if (current !== null) {
+      let prev = new Set();
+      if (previous !== null) {
+        prev = new Set(Object.keys(previous));
+      }
+      const curr = new Set(Object.keys(current));
+      // valid if in current but not previous
+      // invalid if in previous but not current
+      for (const key of curr) {
+        // valid so draw
+        if (!prev.has(key)) {
+          // keys_valid.push(key);
+          const json = current[key];
+          const camera = drawCamera(json, key);
+          sceneTree.set_object_from_path([EVAL_CAMERAS_NAME, key], camera);
+        }
+      }
+      for (const key of prev) {
+        // invalid so delete
+        if (!curr.has(key) || current[key] === null) {
+          // keys_invalid.push(key);
+          sceneTree.delete([EVAL_CAMERAS_NAME, key]);
+        }
+      }
+    } else {
+      sceneTree.delete([EVAL_CAMERAS_NAME]);
+    }
+  };
+  subscribe_to_changes(selector_fn_eval_cameras, fn_value_eval_cameras);
+
+  // draw cubes
+  // NOTE: this has some issues right now! it won't
+  // update the camera on an individual change w/o deleting first
+  const selector_fn_cubes = (state) => {
+    return state.sceneState.cubes;
+  };
+  const fn_value_cubes = (previous, current) => {
+    if (current !== null) {
+      let prev = new Set();
+      if (previous !== null) {
+        prev = new Set(Object.keys(previous));
+      }
+      const curr = new Set(Object.keys(current));
+      // valid if in current but not previous
+      // invalid if in previous but not current
+      for (const key of curr) {
+        // valid so draw
+        if (!prev.has(key)) {
+          // keys_valid.push(key);
+          const json = current[key];
+          const line = drawSceneBox(json);
+          sceneTree.set_object_from_path([CUBES_NAME, key], line);
+        }
+      }
+      for (const key of prev) {
+        // invalid so delete
+        if (!curr.has(key) || current[key] === null) {
+          // keys_invalid.push(key);
+          sceneTree.delete([CUBES_NAME, key]);
+        }
+      }
+    } else {
+      sceneTree.delete([CUBES_NAME]);
+    }
+  };
+  subscribe_to_changes(selector_fn_cubes, fn_value_cubes);
+
   // Check for clicks on training cameras
   const mouseVector = new THREE.Vector2();
   const raycaster = new THREE.Raycaster();

nerfstudio/viewer/app/src/reducer.js

@@ -66,6 +66,8 @@ const initialState = {
   sceneState: {
     sceneBox: null,
     cameras: null,
+    eval_cameras: null,
+    cubes: null,
   },
 };
 

nerfstudio/viewer/server/viewer_utils.py

@@ -325,7 +325,7 @@ class ViewerState:
         # draw indices, roughly evenly spaced
         return np.linspace(0, total_num - 1, num_display_images, dtype=np.int32).tolist()
 
-    def init_scene(self, dataset: InputDataset, start_train=True) -> None:
+    def init_scene(self, dataset: InputDataset, start_train=True, eval_dataset: InputDataset = None) -> None:
         """Draw some images and the scene aabb in the viewer.
 
         Args:
@@ -345,6 +345,7 @@ class ViewerState:
         # clear the current scene
         self.vis["sceneState/sceneBox"].delete()
         self.vis["sceneState/cameras"].delete()
+        self.vis["sceneState/eval_cameras"].delete()
 
         # draw the training cameras and images
         image_indices = self._pick_drawn_image_idxs(len(dataset))
@@ -354,6 +355,15 @@ class ViewerState:
             camera_json = dataset.cameras.to_json(camera_idx=idx, image=bgr, max_size=100)
             self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json)
 
+        if eval_dataset is not None:
+            # draw the eval cameras and images
+            image_indices = self._pick_drawn_image_idxs(len(eval_dataset))
+            for idx in image_indices:
+                image = eval_dataset[idx]["image"]
+                bgr = image[..., [2, 1, 0]]
+                camera_json = eval_dataset.cameras.to_json(camera_idx=idx, image=bgr, max_size=100)
+                self.vis[f"sceneState/eval_cameras/{idx:06d}"].write(camera_json)
+
         # draw the scene box (i.e., the bounding box)
         json_ = dataset.scene_box.to_json()
         self.vis["sceneState/sceneBox"].write(json_)

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "appdirs>=1.4",
     "av>=9.2.0",
     "cryptography>=38",
-    "tyro>=0.3.31",
+    "tyro>=0.4.0",
     "gdown>=4.6.0",
     "ninja>=1.10",
     "functorch>=0.2.1",

scripts/process_data.py

@@ -254,8 +254,8 @@ class ProcessVideo:
     3. Calculates the camera poses for each image using `COLMAP <https://colmap.github.io/>`_.
     """
 
-    data: Path
-    """Path the data, either a video file or a directory of images."""
+    data: Union[Tuple[Path], Tuple[Path, Path]]
+    """Data. Either a video file or a directory of images. If two, then store processed images in separate folders."""
     output_dir: Path
     """Path to the output directory."""
     num_frames_target: int = 300
@@ -321,6 +321,7 @@ class ProcessVideo:
         summary_log = []
         # Convert video to images
         if self.camera_type == "equirectangular":
+            assert len(self.data) == 1, "Only one video can be processed at a time for equirectangular camera type"
             # create temp images folder to store the equirect and perspective images
             temp_image_dir = self.output_dir / "temp_images"
             temp_image_dir.mkdir(parents=True, exist_ok=True)
@@ -332,13 +333,31 @@ class ProcessVideo:
                 verbose=self.verbose,
             )
         else:
+            num_frames_target = int(self.num_frames_target / len(self.data))
             summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
-                self.data,
+                self.data[0],
                 image_dir=image_dir,
                 num_frames_target=self.num_frames_target,
                 crop_factor=self.crop_factor,
                 verbose=self.verbose,
             )
+            for i in range(1, len(self.data)):
+                temp_image_dir = self.output_dir / "temp_images"
+                temp_image_dir.mkdir(parents=True, exist_ok=True)
+                summary_log_i, num_extracted_frames_i = process_data_utils.convert_video_to_images(
+                    self.data[i],
+                    image_dir=temp_image_dir,
+                    num_frames_target=num_frames_target,
+                    verbose=self.verbose,
+                )
+                summary_log += summary_log_i
+                num_extracted_frames += num_extracted_frames_i
+                # copy the images from the temp directory to the image directory
+                process_data_utils.copy_images(
+                    temp_image_dir, image_dir=image_dir, verbose=False, image_prefix=f"frame_{i}_", keep_image_dir=True
+                )
+                # remove the temp directory
+                shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True)
 
         # Generate planar projections if equirectangular
         if self.camera_type == "equirectangular":

scripts/render.py

@@ -36,8 +36,9 @@ from nerfstudio.cameras.camera_paths import (
 from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.model_components import renderers
+from nerfstudio.model_components.renderers import DepthRenderer
 from nerfstudio.pipelines.base_pipeline import Pipeline
-from nerfstudio.utils import install_checks
+from nerfstudio.utils import colormaps, install_checks
 from nerfstudio.utils.eval_utils import eval_setup
 from nerfstudio.utils.rich_utils import ItersPerSecColumn
 
@@ -54,6 +55,8 @@ def _render_trajectory_video(
     seconds: float = 5.0,
     output_format: Literal["images", "video"] = "video",
     camera_type: CameraType = CameraType.PERSPECTIVE,
+    depth_near_plane: Optional[float] = None,
+    depth_far_plane: Optional[float] = None,
 ) -> None:
     """Helper function to create a video of the spiral trajectory.
 
@@ -67,12 +70,18 @@ def _render_trajectory_video(
         seconds: Length of output video.
         output_format: How to save output data.
         camera_type: Camera projection format type.
+        depth_near_plane: Closest depth to consider. If None, use min image value.
+        depth_far_plane: Furthest depth to consider. If None, use max image value.
     """
     CONSOLE.print("[bold green]Creating trajectory " + output_format)
     cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
     cameras = cameras.to(pipeline.device)
     fps = len(cameras) / seconds
 
+    # set the pipeline to use median depth
+    # TODO(ethan): don't hardcode this...
+    pipeline.model.renderer_depth = DepthRenderer(method="median")
+
     progress = Progress(
         TextColumn(":movie_camera: Rendering :movie_camera:"),
         BarColumn(),
@@ -115,6 +124,7 @@ def _render_trajectory_video(
                         outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
 
                 render_image = []
+                depth_raw = None
                 for rendered_output_name in rendered_output_names:
                     if rendered_output_name not in outputs:
                         CONSOLE.rule("Error", style="red")
@@ -123,14 +133,52 @@ def _render_trajectory_video(
                             f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center"
                         )
                         sys.exit(1)
+                    # apply colormap to depth outputs
+                    if rendered_output_name == "depth":
+                        depth_raw = outputs[rendered_output_name]
+                        outputs[rendered_output_name] = colormaps.apply_depth_colormap(
+                            outputs[rendered_output_name],
+                            accumulation=outputs["accumulation"],
+                            near_plane=depth_near_plane,
+                            far_plane=depth_far_plane,
+                        )
+                    if rendered_output_name == "visibility_median_count":
+                        # NOTE(ethan): we cap the number of visible images to 255 here
+                        outputs[rendered_output_name] = outputs[rendered_output_name].clamp(0, 255).to(torch.uint8)
+                        if output_format == "video":
+                            # rescale to [0, 1] is needed to concatenate into a video format (see below)
+                            outputs[rendered_output_name] = outputs[rendered_output_name].float() / 255.0
                     output_image = outputs[rendered_output_name].cpu().numpy()
                     if output_image.shape[-1] == 1:
                         output_image = np.concatenate((output_image,) * 3, axis=-1)
                     render_image.append(output_image)
-                render_image = np.concatenate(render_image, axis=1)
+
+                # TODO: check for eval-images mode
+                # if, so then save the GT image as well!
+                if True:
+                    batch = pipeline.datamanager.eval_dataloader.input_dataset[camera_idx]
+                    rgb_gt = batch["image"]
+                    # resize to match the rendered image
+                    h, w = output_image.shape[:2]
+                    rgb_gt = media.resize_image(rgb_gt, (h, w))
+                    render_image = [rgb_gt] + render_image
+                    names = ["rgb_gt"] + rendered_output_names
+                else:
+                    names = rendered_output_names
+
                 if output_format == "images":
-                    media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image)
+                    # for images, we keep outputs in different folders
+                    for name, im in zip(names, render_image):
+                        folder = output_image_dir / name
+                        folder.mkdir(parents=True, exist_ok=True)
+                        media.write_image(folder / f"{camera_idx:05d}.png", im)
+                    if depth_raw is not None:
+                        folder = output_image_dir / "depth_raw"
+                        folder.mkdir(parents=True, exist_ok=True)
+                        np.save(folder / f"{camera_idx:05d}.npy", depth_raw.cpu().numpy())
                 if output_format == "video":
+                    # for videos, we concatenate images horizontally
+                    render_image = np.concatenate(render_image, axis=1)
                     if writer is None:
                         render_width = int(render_image.shape[1])
                         render_height = int(render_image.shape[0])
@@ -263,9 +311,9 @@ class RenderTrajectory:
     """Path to config YAML file."""
     rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
     """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
-    traj: Literal["spiral", "filename", "interpolate"] = "spiral"
+    traj: Literal["spiral", "filename", "interpolate", "eval-images"] = "spiral"
     """Trajectory type to render. Select between spiral-shaped trajectory, trajectory loaded from
-    a viewer-generated file and interpolated camera paths from the eval dataset."""
+    a viewer-generated file, interpolated camera paths from the eval dataset, of the eval cameras."""
     downscale_factor: int = 1
     """Scaling factor to apply to the camera image resolution."""
     camera_path_filename: Path = Path("camera_path.json")
@@ -280,13 +328,17 @@ class RenderTrajectory:
     """Number of interpolation steps between eval dataset cameras."""
     eval_num_rays_per_chunk: Optional[int] = None
     """Specifies number of rays per chunk during eval."""
+    depth_near_plane: Optional[float] = None
+    """Specifies the near plane for depth rendering."""
+    depth_far_plane: Optional[float] = None
+    """Specifies the far plane for depth rendering."""
 
     def main(self) -> None:
         """Main function."""
         _, pipeline, _ = eval_setup(
             self.load_config,
             eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
-            test_mode="test" if self.traj in ["spiral", "interpolate"] else "inference",
+            test_mode="test" if self.traj in ["spiral", "interpolate", "eval-images"] else "inference",
         )
 
         install_checks.check_ffmpeg_installed()
@@ -319,6 +371,10 @@ class RenderTrajectory:
             camera_path = get_interpolated_camera_path(
                 cameras=pipeline.datamanager.eval_dataloader.cameras, steps=self.interpolation_steps
             )
+        elif self.traj == "eval-images":
+            camera_type = CameraType.PERSPECTIVE
+            # TODO(ethan): support any camera type, reading from the cameras object
+            camera_path = pipeline.datamanager.eval_dataloader.cameras
         else:
             assert_never(self.traj)
 
@@ -332,6 +388,8 @@ class RenderTrajectory:
             seconds=seconds,
             output_format=self.output_format,
             camera_type=camera_type,
+            depth_near_plane=self.depth_near_plane,
+            depth_far_plane=self.depth_far_plane,
         )