appearance working

nerfstudio/fields/nerfacto_field.py

@@ -105,6 +105,8 @@ class TCNNNerfactoField(Field):
         num_semantic_classes: int = 100,
         use_pred_normals: bool = False,
         use_average_appearance_embedding: bool = False,
+        use_eval_appearance_embedding: bool = False,
+        num_eval_images: Optional[int] = None,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
         super().__init__()
@@ -114,8 +116,10 @@ class TCNNNerfactoField(Field):
 
         self.spatial_distortion = spatial_distortion
         self.num_images = num_images
+        self.num_eval_images = num_eval_images
         self.appearance_embedding_dim = appearance_embedding_dim
         self.use_average_appearance_embedding = use_average_appearance_embedding
+        self.use_eval_appearance_embedding = use_eval_appearance_embedding
         self.use_appearance_embedding = use_appearance_embedding
         self.use_transient_embedding = use_transient_embedding
         self.use_semantics = use_semantics
@@ -123,6 +127,7 @@ class TCNNNerfactoField(Field):
 
         if self.use_appearance_embedding:
             self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
+            self.embedding_appearance_eval = None
 
         base_res = 16
         features_per_level = 2
@@ -247,6 +252,16 @@ class TCNNNerfactoField(Field):
         density = trunc_exp(density_before_activation.to(positions))
         return density, base_mlp_out
 
+    def initialize_embedding_appearance_eval(self):
+        """Initializes the appearance embeddings for per-image evaluation.
+        We initialize them to the mean of the appearance embeddings for training."""
+        assert self.num_eval_images and self.num_eval_images > 0, "Number of eval images is not set or is 0."
+        device = self.aabb.device
+        self.embedding_appearance_eval = Embedding(self.num_eval_images, self.appearance_embedding_dim)
+        ones = torch.ones_like(self.embedding_appearance_eval.embedding.weight).to(device)
+        mean = self.embedding_appearance.mean(dim=0)
+        self.embedding_appearance_eval.embedding.weight = torch.nn.Parameter(ones * mean)
+
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
         assert density_embedding is not None
         outputs = {}
@@ -263,15 +278,18 @@ class TCNNNerfactoField(Field):
         if self.use_appearance_embedding:
             if self.training:
                 embedded_appearance = self.embedding_appearance(camera_indices)
+            elif self.embedding_appearance_eval:
+                # if eval appearance embedding exists, use it
+                embedded_appearance = self.embedding_appearance_eval(camera_indices)
+            elif self.use_average_appearance_embedding:
+                # otherwise, if enabled, use the average appearance embedding
+                embedded_appearance = torch.ones(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                ) * self.embedding_appearance.mean(dim=0)
             else:
-                if self.use_average_appearance_embedding:
-                    embedded_appearance = torch.ones(
-                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
-                    ) * self.embedding_appearance.mean(dim=0)
-                else:
-                    embedded_appearance = torch.zeros(
-                        (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
-                    )
+                embedded_appearance = torch.zeros(
+                    (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
+                )
 
         # transients
         if self.use_transient_embedding and self.training:

nerfstudio/models/base_model.py

@@ -69,12 +69,14 @@ class Model(nn.Module):
         config: ModelConfig,
         scene_box: SceneBox,
         num_train_data: int,
+        num_eval_data: Optional[int] = None,
         **kwargs,
     ) -> None:
         super().__init__()
         self.config = config
         self.scene_box = scene_box
         self.num_train_data = num_train_data
+        self.num_eval_data = num_eval_data
         self.kwargs = kwargs
         self.collider = None
 

nerfstudio/models/nerfacto.py

@@ -141,6 +141,7 @@ class NerfactoModel(Model):
             log2_hashmap_size=self.config.log2_hashmap_size,
             spatial_distortion=scene_contraction,
             num_images=self.num_train_data,
+            num_eval_images=self.num_eval_data,
             use_pred_normals=self.config.predict_normals,
             use_appearance_embedding=self.config.use_appearance_embedding,
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,

nerfstudio/pipelines/base_pipeline.py

@@ -46,6 +46,7 @@ from nerfstudio.data.datamanagers.base_datamanager import (
     VanillaDataManager,
     VanillaDataManagerConfig,
 )
+from nerfstudio.model_components.losses import MSELoss
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.utils import profiler
@@ -191,7 +192,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
     """Specifies the datamanager config."""
     model: ModelConfig = ModelConfig()
     """Specifies the model config."""
-    eval_optimize_cameras: bool = True
+    eval_optimize_cameras: bool = False
     """Whether to optimize the cameras during evaluation."""
     num_pose_iters: int = 10
     """Number of iterations to optimize the cameras."""
@@ -239,6 +240,7 @@ class VanillaPipeline(Pipeline):
         self._model = config.model.setup(
             scene_box=self.datamanager.train_dataset.scene_box,
             num_train_data=len(self.datamanager.train_dataset),
+            num_eval_data=len(self.datamanager.eval_dataset),
             metadata=self.datamanager.train_dataset.metadata,
         )
         self.model.to(device)
@@ -354,7 +356,7 @@ 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=6e-4, eps=1e-15)
+                    optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=1e-3, eps=1e-15)
 
                     # check that camera optimizer is on
                     if self.config.datamanager.camera_optimizer.mode == "off":
@@ -368,20 +370,19 @@ class VanillaPipeline(Pipeline):
                     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)
+                        # print("pose opt iter: ", i)
                         # optimize the ray generator's camera optimizer
                         optimizer.zero_grad()
-                        print(self.datamanager.eval_camera_optimizer.pose_adjustment)
+                        # 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())
+                        # import mediapy as media
+                        # _, images_dict = self.model.get_image_metrics_and_images(outputs, batch)
+                        # media.write_image("camera_opt.png", images_dict["img"].detach().cpu().numpy())
 
                         loss_dict["rgb_loss"].backward()
                         optimizer.step()
@@ -389,25 +390,50 @@ class VanillaPipeline(Pipeline):
                     self.eval()
                     print("Done optimizing camera poses.")
 
-                # optimize appearance on half image
                 if self.config.eval_optimize_appearance:
+                    print("Optimizing appearance...")
+                    scale_factor = 0.25
+                    self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(scale_factor)
+                    self.train()
 
-                    # get the parameters for the appearance embedding
-                    # optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=6e-4, eps=1e-15)
+                    # initial the parameters as the mean appearance embedding
+                    self.model.field.initialize_embedding_appearance_eval()
+                    appearance_parameters = self.model.field.embedding_appearance_eval.parameters()
+                    optimizer = torch.optim.Adam(appearance_parameters, lr=1e-3, eps=1e-15)
+
+                    side = random.choice([0, 1])
+                    half_width = width // 2
+                    rgb_loss = MSELoss()
 
                     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
+                        # print("appearance iter: ", i)
+                        optimizer.zero_grad()
+                        # print(self.model.field.embedding_appearance_eval.embedding.weight)
+                        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)
+
+                        # compute loss on half the image
+                        image = batch["image"].to(self.device)
+                        if side == 0:
+                            loss = rgb_loss(image[:, :half_width], outputs["rgb"][:, :half_width])
+                        else:
+                            loss = rgb_loss(image[:, half_width:], outputs["rgb"][:, half_width:])
+
+                        # save the image
+                        # import mediapy as media
+                        # _, images_dict = self.model.get_image_metrics_and_images(outputs, batch)
+                        # media.write_image("appearance_opt.png", images_dict["img"].detach().cpu().numpy())
+
+                        loss.backward()
+                        optimizer.step()
+
+                    self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(1.0 / scale_factor)
+                    # delete the appearance embedding
+                    self.model.field.embedding_appearance_eval = None
+                    self.eval()
+                    print("Done optimizing appearance.")
+
                 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)