nerfacto experiments code

This commit is contained in:
Ethan Weber committed 2023-01-09 23:19:05 +00:00
1 parent 8fd7ea13fe
commit 2f3a7e82c9
4 files changed
+326 -55

No files matched your search

+93 -46
View File
@@ -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)
+195
View File
@@ -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
}
+20
View File
@@ -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")],