From 5447b219aa35522fd4b3a93588fb510627ec00aa Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Tue, 25 Jul 2023 01:41:05 +0000 Subject: [PATCH] changes before merging main --- nerfstudio/configs/method_configs.py | 35 ++++++++++ nerfstudio/pipelines/base_pipeline.py | 93 +++++++++++++++++++++++---- 2 files changed, 115 insertions(+), 13 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 027b01352..219f3d4fb 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -73,6 +73,7 @@ method_configs: Dict[str, TrainerConfig] = {} descriptions = { "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.", "depth-nerfacto": "Nerfacto with depth supervision.", + "visibility-nerfacto": "Nerfacto with visibility loss turned on.", "volinga": "Real-time rendering model from Volinga. Directly exportable to NVOL format at https://volinga.ai/", "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.", "instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes", @@ -194,6 +195,40 @@ method_configs["depth-nerfacto"] = TrainerConfig( vis="viewer", ) +method_configs["visibility-nerfacto"] = TrainerConfig( + method_name="visibility-nerfacto", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", + optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), + ), + ), + model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), + use_visibility_loss=True, + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + method_configs["volinga"] = TrainerConfig( method_name="volinga", steps_per_eval_batch=500, diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 1e1e0fab5..9f9ff2de3 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -25,24 +25,16 @@ from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Uni import torch import torch.distributed as dist -from rich.progress import ( - BarColumn, - MofNCompleteColumn, - Progress, - TextColumn, - TimeElapsedColumn, -) +from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn from torch import nn +from torch.cuda.amp.grad_scaler import GradScaler from torch.nn import Parameter from torch.nn.parallel import DistributedDataParallel as DDP -from torch.cuda.amp.grad_scaler import GradScaler +from nerfstudio.fields.visibility_field import VisibilityField + from nerfstudio.configs import base_config as cfg -from nerfstudio.data.datamanagers.base_datamanager import ( - DataManager, - VanillaDataManager, - VanillaDataManagerConfig, -) +from nerfstudio.data.datamanagers.base_datamanager import DataManager, VanillaDataManager, VanillaDataManagerConfig from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import profiler @@ -209,6 +201,14 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" + use_visibility_loss: bool = False + """whether to apply the visibility loss""" + visibility_loss_quantity: Literal["weights", "densities"] = "densities" + """whether to apply the visibility loss to weights or densities""" + visibility_loss_mult: float = 1.0 + """multiplier for the visibility loss""" + visibility_min_views: int = 1 + """minimum number of views for the visibility loss""" class VanillaPipeline(Pipeline): @@ -263,6 +263,15 @@ class VanillaPipeline(Pipeline): self._model = typing.cast(Model, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True)) dist.barrier(device_ids=[local_rank]) + if self.config.use_visibility_loss: + # bounding box of the scene + self.aabb = self.datamanager.train_dataset.scene_box.aabb.to(self.device) * self.config.aabb_scalar + + # initialize the visibility field + cameras = self.datamanager.train_dataparser_outputs.cameras.to(self.device) + # TODO(ethan): make sure only the training cameras are being used + self.model.visibility_field = VisibilityField(cameras).to(self.device) + @property def device(self): """Returns the device that the model is on.""" @@ -294,6 +303,64 @@ class VanillaPipeline(Pipeline): loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) + if self.config.use_visibility_loss: + # randomly sample a center within the aabb + center = (torch.rand(3, device=self.device) * (self.aabb[1] - self.aabb[0]) + self.aabb[0]).tolist() + cameras, _, _ = random_train_pose( + size=self.config.visibility_num_rays, + resolution=1, + device=self.device, + radius_mean=self.config.aabb_scalar * math.sqrt(3.0), + radius_std=0.0, + central_rotation_range=(0, 360), + vertical_rotation_range=(-90, 90), + focal_range=self.config.focal_range, + jitter_std=0, + center=center, + ) + # We only use the first camera index, but it doesn't matter since we don't care about appearance embeddings + # when using the visibility loss. + camera_indices = torch.tensor([0] * self.config.visibility_num_rays).unsqueeze(-1) + ray_bundle_rays = cameras.generate_rays(camera_indices) + # ray_bundle is (1, 1, self.config.visibility_num_rays) + ray_bundle_rays = ray_bundle_rays.flatten() + model_outputs_rays = self.model(ray_bundle_rays) + quantity_list = model_outputs_rays[f"{self.config.visibility_loss_quantity}_list"] # weights or densities + visibility_loss = 0.0 + for i in range(len(quantity_list)): + quantity_samples = quantity_list[i] + ray_samples = model_outputs_rays["ray_samples_list"][i] + visibility_samples = self.model.visibility_field(ray_samples) + + with torch.no_grad(): + if ( + self.config.sample_method in ["importance", "random"] + and self.config.weight_grid_quantity == "visibility" + ): + if step % self.config.weight_grid_update_per_step == 0: + if i - 1 == len(quantity_list): + # update the weight grid + # get positions + positions = ray_samples.frustums.get_positions() + print(positions.shape) + normalized_positions = SceneBox.get_normalized_positions(positions, self.aabb) + # UPDATE THE WEIGHT GRID + normalized_positions = normalized_positions.view(-1, 3) + quant = visibility_samples.view(-1, 1) + # only use positions between 0 inclusive and 1 exclusive + mask = (normalized_positions >= 0.0) & (normalized_positions < 1.0) + mask = mask[:, 0] & mask[:, 1] & mask[:, 2] + normalized_positions = normalized_positions[mask] + quant = quant[mask] + if normalized_positions.numel() != 0: + self.weight_grid.update(normalized_positions, quant) + + quantity_samples_masked = quantity_samples[visibility_samples < self.config.visibility_min_views] + if quantity_samples_masked.numel() != 0: + visibility_loss_i = self.config.visibility_loss_mult * quantity_samples_masked.mean() + visibility_loss += visibility_loss_i + loss_dict["visibility_loss"] = visibility_loss + return model_outputs, loss_dict, metrics_dict def forward(self):