diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index b754ae206..7aa873aee 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -406,13 +406,25 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ray_bundle = self.eval_ray_generator(ray_indices) return ray_bundle, batch - def get_eval_image(self, idx: int) -> Tuple[RayBundle, Dict]: - """Returns the full image at the given index.""" + def get_eval_image(self, idx: int, scale_factor: float = 1.0) -> Tuple[RayBundle, Dict]: + """Returns the full image at the given index. + + Args: + idx: index of the image to get + scale_factor: scale factor to apply to the image + """ image_batch = self.eval_dataset[idx] image_batch = self.config.collate_fn([image_batch]) image_batch = get_dict_to_torch(image_batch, device=self.device, exclude=["image"]) # image_batch["image"].shape == (1, H, W, 3) # this will sample the entire image + # we downsample the image so it fits in memory. TODO: make this a parameter + print(image_batch["image"].shape) + image_batch["image"] = torch.nn.functional.interpolate( + image_batch["image"].permute(0, 3, 1, 2), scale_factor=scale_factor, mode="bilinear", align_corners=False + ).permute(0, 2, 3, 1) + print(image_batch["image"].shape) + batch = self.eval_pixel_sampler.sample(image_batch, sample_all_pixels=True) ray_indices = batch["indices"] ray_indices_shape = ray_indices.shape # (1, H, W, 3) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index af0d78d96..73de83ca3 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,7 +17,7 @@ Code for sampling pixels. """ import random -from typing import Dict +from typing import Dict, Optional import torch @@ -165,7 +165,9 @@ class PixelSampler: # pylint: disable=too-few-public-methods """ self.num_rays_per_batch = num_rays_per_batch - def sample(self, image_batch: Dict, sample_all_pixels: bool = False) -> Dict: + def sample( + self, image_batch: Dict, num_rays_per_batch: Optional[int] = None, sample_all_pixels: bool = False + ) -> Dict: """Sample an image batch and return a pixel batch. Args: @@ -175,11 +177,15 @@ class PixelSampler: # pylint: disable=too-few-public-methods if isinstance(image_batch["image"], list): image_batch = dict(image_batch.items()) # copy the dictioary so we don't modify the original pixel_batch = collate_image_dataset_batch_list( - image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels + image_batch, + self.num_rays_per_batch if not num_rays_per_batch else num_rays_per_batch, + sample_all_pixels=sample_all_pixels, ) elif isinstance(image_batch["image"], torch.Tensor): pixel_batch = collate_image_dataset_batch( - image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels + image_batch, + self.num_rays_per_batch if not num_rays_per_batch else num_rays_per_batch, + sample_all_pixels=sample_all_pixels, ) else: raise ValueError("image_batch['image'] must be a list or torch.Tensor") diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 6be459343..354558d59 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -158,7 +158,6 @@ class Model(nn.Module): metrics_dict: dictionary of metrics, some of which we can use for loss """ - @torch.no_grad() def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 6df0013c7..61fe45b39 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -281,7 +281,7 @@ class NerfactoModel(Model): metrics_dict = {} image = batch["image"].to(self.device) metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) - if self.training: + if self.training and ("weights_list" in outputs and "ray_samples_list" in outputs): metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"]) return metrics_dict @@ -289,7 +289,7 @@ class NerfactoModel(Model): loss_dict = {} image = batch["image"].to(self.device) loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) - if self.training: + if self.training and ("weights_list" in outputs and "ray_samples_list" in outputs): loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( outputs["weights_list"], outputs["ray_samples_list"] ) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 81f2dc3bc..802b60c01 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -37,6 +37,8 @@ from torch import nn from torch.nn import Parameter from torch.nn.parallel import DistributedDataParallel as DDP from typing_extensions import Literal +from rich.progress import track +from tqdm import tqdm from nerfstudio.configs import base_config as cfg from nerfstudio.data.datamanagers.base_datamanager import ( @@ -189,10 +191,14 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """Specifies the datamanager config.""" model: ModelConfig = ModelConfig() """Specifies the model config.""" - eval_optimize_cameras: bool = False + eval_optimize_cameras: bool = True """Whether to optimize the cameras during evaluation.""" + num_pose_iters: int = 10 + """Number of iterations to optimize the cameras.""" eval_optimize_appearance: bool = False """Whether to optimize the appearance during evaluation.""" + num_appearance_iters: int = 10 + """Number of iterations to optimize the appearance.""" class VanillaPipeline(Pipeline): @@ -307,8 +313,9 @@ class VanillaPipeline(Pipeline): """ self.eval() image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step) - outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch) + with torch.no_grad(): + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch) assert "image_idx" not in metrics_dict metrics_dict["image_idx"] = image_idx assert "num_rays" not in metrics_dict @@ -347,43 +354,63 @@ class VanillaPipeline(Pipeline): # get the eval camera optimizer's parameters camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + "_eval" param_groups = self.datamanager.get_param_groups() - optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=0.1) + optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=6e-4, eps=1e-15) # check that camera optimizer is on if self.config.datamanager.camera_optimizer.mode == "off": raise ValueError( "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses." ) - # TODO: make this a parameter - num_pose_iters = 100 - for i in range(num_pose_iters): - print("pose iter: ", i) - # need to optimize the ray generator's camera optimizer - # TODO: need to add require_grad in various locations to make it work + + print("Optimizing camera poses...") + # rescale the image to 1/4 resolution + scale_factor = 0.25 + self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(scale_factor) + self.train() + for i in range(self.config.num_pose_iters): + print("pose opt iter: ", i) + # optimize the ray generator's camera optimizer optimizer.zero_grad() - camera_ray_bundle, batch = self.datamanager.get_eval_image(idx) + print(self.datamanager.eval_camera_optimizer.pose_adjustment) + camera_ray_bundle, batch = self.datamanager.get_eval_image(idx, scale_factor=0.25) outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) metrics_dict = self.model.get_metrics_dict(outputs, batch) loss_dict = self.model.get_loss_dict(outputs, batch, metrics_dict) + _, images_dict = self.model.get_image_metrics_and_images(outputs, batch) + + # save the image + import mediapy as media + + media.write_image("camera_opt.png", images_dict["img"].detach().cpu().numpy()) + loss_dict["rgb_loss"].backward() optimizer.step() + self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(1.0 / scale_factor) + self.eval() + print("Done optimizing camera poses.") # optimize appearance on half image if self.config.eval_optimize_appearance: - # TODO: make this a parameter - num_appearance_iters = 100 - side = random.choice([0, 1]) - half_width = width // 2 - if side == 0: - half_camera_ray_bundle = camera_ray_bundle[:, :half_width] - else: - half_camera_ray_bundle = camera_ray_bundle[:, half_width:] - for i in range(num_appearance_iters): - # need to optimize the "embedding_appearance" in the nerfacto field - pass - outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch) + # get the parameters for the appearance embedding + # optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=6e-4, eps=1e-15) + + for i in range(self.config.num_appearance_iters): + print("appearance iter: ", i) + # need to optimize the "embedding_appearance" in the nerfacto field + # pass + # side = random.choice([0, 1]) + # half_width = width // 2 + # if side == 0: + # half_camera_ray_bundle = camera_ray_bundle[:, :half_width] + # else: + # half_camera_ray_bundle = camera_ray_bundle[:, half_width:] + # for i in range(num_appearance_iters): + # # need to optimize the "embedding_appearance" in the nerfacto field + # pass + with torch.no_grad(): + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch) assert "num_rays_per_sec" not in metrics_dict metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start) fps_str = "fps"