From 2f3a7e82c912f502063fa368bd8de2fa2c8d00fe Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Mon, 9 Jan 2023 23:19:05 +0000 Subject: [PATCH] nerfacto experiments code --- nerfstudio/models/nerfacto.py | 27 ++- scripts/benchmarking/benchmark_nerfacto.py | 139 ++++++++++----- scripts/benchmarking/create_tables.ipynb | 195 +++++++++++++++++++++ scripts/downloads/download_data.py | 20 +++ 4 files changed, 326 insertions(+), 55 deletions(-) create mode 100644 scripts/benchmarking/create_tables.ipynb diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index b531cff54..a75fb8090 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -53,9 +53,9 @@ from nerfstudio.model_components.renderers import ( NormalsRenderer, RGBRenderer, ) -from nerfstudio.model_components.scene_colliders import NearFarCollider +from nerfstudio.model_components.scene_colliders import NearFarCollider, AABBBoxCollider from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps +from nerfstudio.utils import colormaps, colors @dataclass @@ -67,6 +67,8 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to start sampling.""" far_plane: float = 1000.0 """How far along the ray to stop sampling.""" + use_bounded: bool = False + """Whether the sampling should be bounded inside the cube. Overrides near and far planes.""" background_color: Literal["random", "last_sample"] = "last_sample" """Whether to randomize the background color.""" num_levels: int = 16 @@ -137,13 +139,16 @@ class NerfactoModel(Model): # TODO(ethan): flag to enable scene contraction or not - if self.config.scene_contraction_norm == "inf": - order = float("inf") - elif self.config.scene_contraction_norm == "l2": - order = None + if self.config.use_scene_contraction: + if self.config.scene_contraction_norm == "inf": + order = float("inf") + elif self.config.scene_contraction_norm == "l2": + order = None + else: + raise ValueError("Invalid scene contraction norm") + scene_contraction = SceneContraction(order=order) else: - raise ValueError("Invalid scene contraction norm") - scene_contraction = SceneContraction(order=order) + scene_contraction = None # Fields self.field = TCNNNerfactoField( @@ -196,7 +201,11 @@ class NerfactoModel(Model): ) # Collider - self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + if not self.config.use_bounded: + self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane) + else: + # this will only sample inside the the bounding box + self.collider = AABBBoxCollider(scene_box=self.scene_box) # renderers self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) diff --git a/scripts/benchmarking/benchmark_nerfacto.py b/scripts/benchmarking/benchmark_nerfacto.py index e450217f0..56d8fda6e 100644 --- a/scripts/benchmarking/benchmark_nerfacto.py +++ b/scripts/benchmarking/benchmark_nerfacto.py @@ -1,15 +1,35 @@ """ -Process the Mipnerf360 dataset. +Benchmarking script for nerfstudio paper. + +- nerfacto and instant-ngp methods on mipnerf360 data +- nerfacto ablations """ from nerfstudio.utils.scripts import run_command import GPUtil +from pathlib import Path import time import threading +mipnerf360_capture_names = ["bicycle", "garden", "stump", "room", "counter", "kitchen", "bonsai"] # 7 splits +mipnerf360_table_rows = [ + # nerfacto method + ( + "nerfacto-w/o-pose-app", + "nerfacto", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False nerfstudio-data --downscale-factor 4", + ), + # TODO: instant-ngp method + # ( + # "instant-ngp-w/o-pose-app", + # "instant-ngp", + # "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False nerfstudio-data --downscale-factor 4", + # ), +] -capture_names = [ + +ablations_capture_names = [ "Egypt", "person", "kitchen", @@ -21,88 +41,115 @@ capture_names = [ "sculpture", "Giannini-Hall", ] -table_rows = [ - ("nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True"), + +ablations_table_rows = [ + ("nerfacto", "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True"), ( "w/o-pose", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.datamanager.camera-optimizer.mode off", ), ( "w/o-app", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.model.use-appearance-embedding False", ), ( "w/o-pose-app", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False", ), ( "1-prop-network", + "nerfacto", '--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.num-proposal-samples-per-ray "256" --pipeline.model.num_proposal_iterations 1', ), ( "l2-contraction", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.scene-contraction-norm l2", ), ( "shared-prop-network", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-same-proposal-network True", ), ( "random-background-color", + "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.background-color random", ), + ( + "no-contraction", + "nerfacto", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-bounded True --pipeline.model.use-scene-contraction False nerfstudio-data --scale_factor 0.125", + ), + # experiment with synthetic dataset settings + ( + "synthetic-on-real", + "nerfacto", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False --pipeline.model.use-bounded True --pipeline.model.use-scene-contraction False nerfstudio-data --scale_factor 0.125", + ), ] -# 30K iterations -# make a list of all the jobs that need to be fun -jobs = [] -for capture_name in capture_names: - for table_row_name, table_row_command in table_rows: - command = " ".join( - ( - f"ns-train nerfacto", - "--vis wandb", - f"--data data/nerfstudio/{capture_name}", - "--output-dir outputs/nerfacto-ablations", - "--trainer.steps-per-eval-batch 0 --trainer.steps-per-eval-image 0", - "--trainer.steps-per-eval-all-images 5000 --trainer.max-num-iterations 30001", - f"--wandb-name {capture_name}_{table_row_name}", - f"--experiment-name {capture_name}_{table_row_name}", - table_row_command, +def main(capture_names, table_rows, data_path: Path = Path("data/nerfstudio")): + """Main method.""" + # 30K iterations + + # make a list of all the jobs that need to be fun + jobs = [] + for capture_name in capture_names: + for table_row_name, method, table_row_command in table_rows: + command = " ".join( + ( + f"ns-train {method}", + "--vis wandb", + f"--data { data_path / capture_name}", + "--output-dir outputs/nerfacto-ablations", + "--trainer.steps-per-eval-batch 0 --trainer.steps-per-eval-image 0", + "--trainer.steps-per-eval-all-images 5000 --trainer.max-num-iterations 30001", + f"--wandb-name {capture_name}_{table_row_name}", + f"--experiment-name {capture_name}_{table_row_name}", + table_row_command, + ) ) - ) - jobs.append(command) + jobs.append(command) -while jobs: + while jobs: - # check which GPUs have capacity to run these jobs - """Returns the available GPUs.""" - gpu_devices_available = GPUtil.getAvailable(order="first", limit=10, maxMemory=0.1) + # check which GPUs have capacity to run these jobs + """Returns the available GPUs.""" + gpu_devices_available = GPUtil.getAvailable(order="first", limit=10, maxMemory=0.1) - print("Available GPUs: ", gpu_devices_available) + print("Available GPUs: ", gpu_devices_available) - # thread list - threads = [] - while gpu_devices_available and jobs: - gpu = gpu_devices_available.pop(0) - command = f"CUDA_VISIBLE_DEVICES={gpu} " + jobs.pop(0) + # thread list + threads = [] + while gpu_devices_available and jobs: + gpu = gpu_devices_available.pop(0) + command = f"CUDA_VISIBLE_DEVICES={gpu} " + jobs.pop(0) - def task(): - print("Starting command: ", command) - out = run_command(command, verbose=False) - # time.sleep(5) - print("Finished command: ", command) + def task(): + print("Starting command: ", command) + out = run_command(command, verbose=False) + # time.sleep(5) + print("Finished command: ", command) - threads.append(threading.Thread(target=task)) - threads[-1].start() + threads.append(threading.Thread(target=task)) + threads[-1].start() - # NOTE(ethan): need a delay otherwise the wandb/tensorboard naming is messed up - # not sure why? - time.sleep(5) + # NOTE(ethan): need a delay otherwise the wandb/tensorboard naming is messed up + # not sure why? + time.sleep(5) - # wait for all threads to finish - for t in threads: - t.join() + # wait for all threads to finish + for t in threads: + t.join() - print("Finished all threads") + print("Finished all threads") + + +if __name__ == "__main__": + main(mipnerf360_capture_names, mipnerf360_table_rows, data_path=Path("data/nerfstudio-data-mipnerf360")) + # main(ablations_capture_names, ablations_table_rows) diff --git a/scripts/benchmarking/create_tables.ipynb b/scripts/benchmarking/create_tables.ipynb new file mode 100644 index 000000000..2ad5ad562 --- /dev/null +++ b/scripts/benchmarking/create_tables.ipynb @@ -0,0 +1,195 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "import copy\n", + "from benchmark_nerfacto import table_rows\n", + "\n", + "# methods sorted in the same order we want in the paper tables\n", + "methods = [row[0] for row in table_rows]" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "def csv_filename_to_results(csv_filename, num_decimals=2):\n", + " # code to read the csv file\n", + " # and return a dictionary\n", + "\n", + " df = pd.read_csv(csv_filename)\n", + "\n", + " column_titles = list(df.columns.values)\n", + " last_row = list(df.values[-1])\n", + "\n", + " # print(len(column_titles))\n", + " # print(len(last_row))\n", + "\n", + " from collections import defaultdict\n", + " results = defaultdict(dict)\n", + " for column_title, last_row_value in zip(column_titles, last_row):\n", + " if column_title == \"Step\" or \"__MAX\" in column_title or \"__MIN\" in column_title:\n", + " continue\n", + " name = column_title.split(\" - \")[0]\n", + " capture_name, method = name.split(\"_\")\n", + " metric = float(last_row_value)\n", + " # format to have 2 decimal places\n", + " if num_decimals == 2:\n", + " metric = f\"{metric:.2f}\"\n", + " elif num_decimals == 3:\n", + " metric = f\"{metric:.3f}\"\n", + " results[capture_name][method] = metric\n", + "\n", + " # import pprint\n", + " # pprint.pprint(results)\n", + " # print(len(results))\n", + "\n", + " results = dict(results)\n", + " return results" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "# the csv filenames\n", + "# these are generated from exporting CSVs from wandb graphs\n", + "# csv_filename_psrn = \"/home/ethanweber/nerfactory/wandb-psnr.csv\"\n", + "# csv_filename_ssim = \"/home/ethanweber/nerfactory/wandb-ssim.csv\"\n", + "# csv_filename_lpips = \"/home/ethanweber/nerfactory/wandb-lpips.csv\"\n", + "csv_filename_psrn = \"/home/ethanweber/nerfactory/wandb-jan9-psnr.csv\"\n", + "csv_filename_ssim = \"/home/ethanweber/nerfactory/wandb-jan9-ssim.csv\"\n", + "csv_filename_lpips = \"/home/ethanweber/nerfactory/wandb-jan9-lpips.csv\"\n", + "\n", + "results_psrn = csv_filename_to_results(csv_filename_psrn)\n", + "results_ssim = csv_filename_to_results(csv_filename_ssim, num_decimals=3)\n", + "results_lpips = csv_filename_to_results(csv_filename_lpips, num_decimals=3)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_row(method, capture_names):\n", + " # remove spaces from method\n", + " method_text = method.replace(\"-\", \" \")\n", + " string = f\"{method_text}\"\n", + " for idx, capture_name in enumerate(capture_names):\n", + " psnr = results_psrn[capture_name][method]\n", + " ssim = results_ssim[capture_name][method]\n", + " lpips = results_lpips[capture_name][method]\n", + " if idx == len(capture_names) - 1:\n", + " string += \" & \\\\multicolumn{1}{r|}{\" + f\"{psnr} / {ssim} / {lpips}\" + \"}\"\n", + " else:\n", + " string += \" & \\\\multicolumn{1}{r}{\" + f\"{psnr} / {ssim} / {lpips}\" + \"}\"\n", + " return string" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "['Egypt', 'person', 'kitchen', 'plane', 'dozer']\n", + "no contraction & \\multicolumn{1}{r}{18.06 / 0.520 / 0.151} & \\multicolumn{1}{r}{21.96 / 0.511 / 0.125} & \\multicolumn{1}{r}{19.35 / 0.764 / 0.078} & \\multicolumn{1}{r}{18.45 / 0.519 / 0.161} & \\multicolumn{1}{r|}{19.53 / 0.786 / 0.081} \\\\ \n", + "synthetic on real & \\multicolumn{1}{r}{18.51 / 0.546 / 0.119} & \\multicolumn{1}{r}{24.68 / 0.667 / 0.099} & \\multicolumn{1}{r}{20.98 / 0.745 / 0.074} & \\multicolumn{1}{r}{19.32 / 0.578 / 0.102} & \\multicolumn{1}{r|}{21.41 / 0.758 / 0.080} \\\\ \\hline\n", + "\n", + "['floating-tree', 'aspen', 'stump', 'sculpture', 'Giannini-Hall']\n", + "no contraction & \\multicolumn{1}{r}{18.98 / 0.679 / 0.088} & \\multicolumn{1}{r}{15.80 / 0.330 / 0.342} & \\multicolumn{1}{r}{20.69 / 0.704 / 0.075} & \\multicolumn{1}{r}{21.43 / 0.720 / 0.073} & \\multicolumn{1}{r|}{17.75 / 0.596 / 0.134} \\\\ \n", + "synthetic on real & \\multicolumn{1}{r}{21.13 / 0.726 / 0.072} & \\multicolumn{1}{r}{15.53 / 0.301 / 0.332} & \\multicolumn{1}{r}{23.86 / 0.786 / 0.053} & \\multicolumn{1}{r}{22.62 / 0.706 / 0.076} & \\multicolumn{1}{r|}{18.41 / 0.576 / 0.133} \\\\ \\hline\n" + ] + } + ], + "source": [ + "capture_names_row1 = [\n", + " \"Egypt\",\n", + " \"person\",\n", + " \"kitchen\",\n", + " \"plane\",\n", + " \"dozer\"\n", + "]\n", + "capture_names_row2 = [\n", + " \"floating-tree\",\n", + " \"aspen\",\n", + " \"stump\",\n", + " \"sculpture\",\n", + " \"Giannini-Hall\"\n", + "]\n", + "\n", + "print(capture_names_row1)\n", + "for method in methods[:-1]:\n", + " # print(generate_row(method, capture_names_row1) + ' \\\\\\ \\cline{1-1}')\n", + " print(generate_row(method, capture_names_row1) + ' \\\\\\ ')\n", + "print(generate_row(methods[-1], capture_names_row1) + ' \\\\\\ \\hline')\n", + "\n", + "print()\n", + "\n", + "print(capture_names_row2)\n", + "for method in methods[:-1]:\n", + " # print(generate_row(method, capture_names_row2) + ' \\\\\\ \\cline{1-1}')\n", + " print(generate_row(method, capture_names_row2) + ' \\\\\\ ')\n", + "print(generate_row(methods[-1], capture_names_row2) + ' \\\\\\ \\hline')\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "nerfactory", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.6" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "b7b86c17c0dfd6b313529972d424f2e075cc198a46cda005757126afe8133c7b" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 9636d0b58..b5792b6ae 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -79,6 +79,25 @@ class FriendsDownload(DatasetDownload): os.remove(download_path) +@dataclass +class Mipnerf360Download(DatasetDownload): + """Download the friends dataset.""" + + def download(self, save_dir: Path): + """Download the friends dataset.""" + + # https://drive.google.com/file/d/1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5/view?usp=sharing + friends_file_id = "1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5" + + # Download the files + url = "https://data.nerf.studio/nerfstudio-data-mipnerf360.zip" + download_path = str(save_dir / "nerfstudio-data-mipnerf360.zip") + gdown.download(url, output=download_path) + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(str(save_dir)) + os.remove(download_path) + + def grab_file_id(zip_url: str) -> str: """Get the file id from the google drive zip url.""" s = zip_url.split("/d/")[1] @@ -306,6 +325,7 @@ class PhototourismDownload(DatasetDownload): Commands = Union[ Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], Annotated[FriendsDownload, tyro.conf.subcommand(name="friends")], + Annotated[Mipnerf360Download, tyro.conf.subcommand(name="mipnerf360")], Annotated[NerfstudioDownload, tyro.conf.subcommand(name="nerfstudio")], Annotated[Record3dDownload, tyro.conf.subcommand(name="record3d")], Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")],