diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 88b512789..ba46e9c4b 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -32,6 +32,7 @@ from rich.progress import ( Progress, TextColumn, TimeElapsedColumn, + Console ) from torch import nn from torch.nn import Parameter @@ -51,6 +52,8 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import profiler +CONSOLE = Console(width=120) + def module_wrapper(ddp_or_model: Union[DDP, Model]) -> Model: """ @@ -194,12 +197,14 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """Specifies the model config.""" eval_optimize_cameras: bool = False """Whether to optimize the cameras during evaluation.""" - num_pose_iters: int = 10 + eval_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 + eval_num_appearance_iters: int = 10 """Number of iterations to optimize the appearance.""" + eval_image_scale_factor: float = 0.25 + """Scale factor to use for evaluation images so that they fit in memory.""" class VanillaPipeline(Pipeline): @@ -353,6 +358,8 @@ class VanillaPipeline(Pipeline): # optimize camera poses on full image if self.config.eval_optimize_cameras: + # raise error if the method not nerfacto + # get the eval camera optimizer's parameters camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + "_eval" # reinitialize the parameters @@ -366,17 +373,18 @@ class VanillaPipeline(Pipeline): "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses." ) - print("Optimizing camera poses...") + # CONSOLE.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.datamanager.eval_ray_generator.cameras.rescale_output_resolution( + self.config.eval_image_scale_factor + ) self.train() - for i in range(self.config.num_pose_iters): - # print("pose opt iter: ", i) + for _ in range(self.config.eval_num_pose_iters): # optimize the ray generator's camera optimizer optimizer.zero_grad() - # print(self.datamanager.eval_camera_optimizer.pose_adjustment) - camera_ray_bundle, batch = self.datamanager.get_eval_image(idx, scale_factor=0.25) + camera_ray_bundle, batch = self.datamanager.get_eval_image( + idx, scale_factor=self.config.eval_image_scale_factor + ) 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) @@ -388,14 +396,17 @@ class VanillaPipeline(Pipeline): loss_dict["rgb_loss"].backward() optimizer.step() - self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(1.0 / scale_factor) + self.datamanager.eval_ray_generator.cameras.rescale_output_resolution( + 1.0 / self.config.eval_image_scale_factor + ) self.eval() - print("Done optimizing camera poses.") + # CONSOLE.print("Done optimizing camera poses.") if self.config.eval_optimize_appearance: - print("Optimizing appearance...") - scale_factor = 0.25 - self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(scale_factor) + # CONSOLE.print("Optimizing appearance...") + self.datamanager.eval_ray_generator.cameras.rescale_output_resolution( + self.config.eval_image_scale_factor + ) self.train() # initial the parameters as the mean appearance embedding @@ -407,11 +418,11 @@ class VanillaPipeline(Pipeline): half_width = width // 2 rgb_loss = MSELoss() - for i in range(self.config.num_appearance_iters): - # print("appearance iter: ", i) + for _ in range(self.config.eval_num_appearance_iters): 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) + camera_ray_bundle, batch = self.datamanager.get_eval_image( + idx, scale_factor=self.config.eval_image_scale_factor + ) outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) metrics_dict = self.model.get_metrics_dict(outputs, batch) @@ -430,11 +441,11 @@ class VanillaPipeline(Pipeline): loss.backward() optimizer.step() - self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(1.0 / scale_factor) + self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(1.0 / self.config.eval_image_scale_factor) # delete the appearance embedding self.model.field.embedding_appearance_eval = None self.eval() - print("Done optimizing appearance.") + # CONSOLE.print("Done optimizing appearance.") with torch.no_grad(): outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) diff --git a/scripts/benchmarking/benchmark_nerfacto.py b/scripts/benchmarking/benchmark_nerfacto.py new file mode 100644 index 000000000..082590cad --- /dev/null +++ b/scripts/benchmarking/benchmark_nerfacto.py @@ -0,0 +1,52 @@ +""" +Process the Mipnerf360 dataset. +""" + +from nerfstudio.utils.scripts import run_command + +capture_names = [ + "Egypt", + "person", + "kitchen", + "plane", + "dozer", + "floating-tree", + "aspen", + "stump", + "sculpture", + "Giannini-Hall", +] +table_rows = [ + ("nerfacto", ""), + ("w/o-pose", "--pipeline.datamanager.camera-optimizer.mode off"), + ("w/o-app", "--pipeline.model.use-appearance-embedding False"), + ( + "w/o-pose-app", + "--pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False", + ), + ( + "1-prop-network", + "--pipeline.model.num-proposal-samples-per-ray \"256\" --pipeline.model.num_proposal_iterations 1", + ), +] + +# 30K iterations + +for capture_name in capture_names: + for table_row_name, table_row_command in table_rows: + + # TODO: find an available GPU + gpu = 0 + + command = " ".join( + ( + # f"CUDA_VISIBLE_DEVICES={gpu}", + f"ns-train nerfacto", + "--vis wandb", + f"--data data/nerfstudio/{capture_name}", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True", + f"--wandb-name {capture_name}-{table_row_name} --experiment-name {capture_name}-{table_row_name}", + table_row_command, + ) + ) + print(command) diff --git a/scripts/datasets/process_mipnerf360.py b/scripts/datasets/process_mipnerf360.py index fbc235eb9..825fc9e02 100644 --- a/scripts/datasets/process_mipnerf360.py +++ b/scripts/datasets/process_mipnerf360.py @@ -91,3 +91,38 @@ with Pool() as p: # DATASET="plane" METHOD="nerfacto-pose" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}" --experiment-name "${METHOD}-${DATASET}" --trainer.load-dir outputs/nerfacto-no-pose-plane/nerfacto/2022-12-16_220419/nerfstudio_models --trainer.steps-per-eval-all-images 10 --pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False # DATASET="plane" METHOD="nerfacto-pose-app" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}" --experiment-name "${METHOD}-${DATASET}" --trainer.load-dir outputs/nerfacto-no-pose-plane/nerfacto/2022-12-16_220419/nerfstudio_models --trainer.steps-per-eval-all-images 10 --pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True # DATASET="plane" METHOD="nerfacto-app" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}" --experiment-name "${METHOD}-${DATASET}" --trainer.load-dir outputs/nerfacto-no-pose-plane/nerfacto/2022-12-16_220419/nerfstudio_models --trainer.steps-per-eval-all-images 10 --pipeline.eval_optimize_cameras False --pipeline.eval_optimize_appearance True + + +capture_names = [ + "Egypt", + "person", + "kitchen", + "plane", + "dozer", + "floating-tree", + "aspen", + "stump", + "sculpture", + "Giannini-Hall", +] +table_row = [ + ("nerfacto", ""), + ("w/o-pose", "--pipeline.datamanager.camera-optimizer.mode off"), + ("w/o-app", "--pipeline.model.use-appearance-embedding False"), + ( + "w/o-pose-app", + "--pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False", + ), + ("1-prop-network", "--pipeline.model.num_proposal_samples_per_ray (256,) --pipeline.model.num_proposal_iterations 1"), +] + +# 30K iterations + +command = """ +CAPTURE_NAME="plane" +TABLE_ROW="nerfacto-app" +&& +ns-train nerfacto --vis wandb --data "data/nerfstudio/${CAPTURE_NAME}" --wandb-name "${DATASET}-${TABLE_ROW}" --experiment-name "${DATASET}-${TABLE_ROW}" +--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True +--trainer.steps-per-eval-all-images 30000 --trainer.max-num-iterations 30001 +"""