From 984f36738a4ff9f51e3bc1dcc39d7d05ec13598c Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 7 Apr 2023 21:27:11 +0000 Subject: [PATCH] one commit --- nerfstudio/cameras/cameras.py | 2 +- nerfstudio/configs/experiment_config.py | 2 + .../data/dataparsers/nerfstudio_dataparser.py | 105 +++++++++++++----- nerfstudio/data/utils/dataloaders.py | 14 +-- nerfstudio/engine/callbacks.py | 2 + nerfstudio/engine/trainer.py | 15 ++- nerfstudio/fields/visibility_field.py | 90 +++++++++++++++ nerfstudio/model_components/ray_samplers.py | 6 +- nerfstudio/model_components/renderers.py | 19 ++-- nerfstudio/models/nerfacto.py | 41 ++++++- nerfstudio/process_data/process_data_utils.py | 26 ++++- nerfstudio/utils/writer.py | 14 ++- .../viewer/app/src/modules/Scene/Scene.jsx | 76 +++++++++++++ nerfstudio/viewer/app/src/reducer.js | 2 + nerfstudio/viewer/server/viewer_utils.py | 12 +- pyproject.toml | 2 +- scripts/process_data.py | 25 ++++- scripts/render.py | 70 +++++++++++- 18 files changed, 448 insertions(+), 75 deletions(-) create mode 100644 nerfstudio/fields/visibility_field.py diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index e0ed08e4f..a56277149 100644 --- a/nerfstudio/cameras/cameras.py +++ b/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) diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 72362befe..c0da16fef 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/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() diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 7b283d747..f5b528369 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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 diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 23d712d7e..93d309dd8 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/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 diff --git a/nerfstudio/engine/callbacks.py b/nerfstudio/engine/callbacks.py index 6be1d8097..8e261c13f 100644 --- a/nerfstudio/engine/callbacks.py +++ b/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): diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 0262703dd..0f9015d45 100644 --- a/nerfstudio/engine/trainer.py +++ b/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: diff --git a/nerfstudio/fields/visibility_field.py b/nerfstudio/fields/visibility_field.py new file mode 100644 index 000000000..855ccea93 --- /dev/null +++ b/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 diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index f5a96255b..145497525 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/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): diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 43f81404c..186f77b29 100644 --- a/nerfstudio/model_components/renderers.py +++ b/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 diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 53173640a..e0176fe74 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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( diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index a977df8b2..0e7df7d86 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/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)) diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index bffb0bd03..56d662770 100644 --- a/nerfstudio/utils/writer.py +++ b/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)) diff --git a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx index ef5a38888..ef7b64e35 100644 --- a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx +++ b/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(); diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js index 21279a547..dd4bd137f 100644 --- a/nerfstudio/viewer/app/src/reducer.js +++ b/nerfstudio/viewer/app/src/reducer.js @@ -66,6 +66,8 @@ const initialState = { sceneState: { sceneBox: null, cameras: null, + eval_cameras: null, + cubes: null, }, }; diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 33fb9afe9..d00db8617 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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_) diff --git a/pyproject.toml b/pyproject.toml index 0c7c338fe..969830d88 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/process_data.py b/scripts/process_data.py index 2eaa9e11c..5977bf850 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -254,8 +254,8 @@ class ProcessVideo: 3. Calculates the camera poses for each image using `COLMAP `_. """ - 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": diff --git a/scripts/render.py b/scripts/render.py index ce2f79dc2..e7c24d0cb 100644 --- a/scripts/render.py +++ b/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, )