diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 1213503eb..ecc4f4202 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/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: diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 354558d59..47b493226 100644 --- a/nerfstudio/models/base_model.py +++ b/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 diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 61fe45b39..2e46b657e 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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, diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 802b60c01..9babe3113 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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)