2 psnr boost w/ camera opt for eval

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)

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")

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.
 

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"]
             )

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"