tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

working on abllations

Browse Source
This commit is contained in:
Ethan Weber 2022-12-23 19:31:02 +00:00
parent d74a0f6ecd
commit 8fd7ea13fe
8 changed files with 156 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -18,6 +18,7 @@ Datamanager.
from __future__ import annotations
import copy
import random
from abc import abstractmethod
from dataclasses import dataclass, field
@ -382,7 +383,10 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
)
self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch)
self.eval_camera_optimizer = self.config.camera_optimizer.setup(
# eval camera optimizer should always have pose opt on
eval_camera_optimizer = copy.deepcopy(self.config.camera_optimizer)
eval_camera_optimizer.mode = "SO3xR3"
self.eval_camera_optimizer = eval_camera_optimizer.setup(
num_cameras=self.eval_dataset.cameras.size, device=self.device
)
self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device), self.eval_camera_optimizer)
@ -450,11 +454,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
camera_opt_params_eval = list(self.eval_camera_optimizer.parameters())
if self.config.camera_optimizer.mode != "off":
assert len(camera_opt_params_train) > 0
assert len(camera_opt_params_eval) > 0
param_groups[self.config.camera_optimizer.param_group + "_train"] = camera_opt_params_train
param_groups[self.config.camera_optimizer.param_group + "_eval"] = camera_opt_params_eval
else:
assert len(camera_opt_params_train) == 0
assert len(camera_opt_params_eval) == 0
assert len(camera_opt_params_eval) > 0
param_groups[self.config.camera_optimizer.param_group + "_eval"] = camera_opt_params_eval
return param_groups

4
nerfstudio/engine/trainer.py
View File

@ -148,10 +148,6 @@ class Trainer:
"optimizer": camera_optimizer_config.optimizer,
"scheduler": camera_optimizer_config.scheduler,
}
# optimizer_config[camera_optimizer_config.param_group + "_eval"] = {
# "optimizer": camera_optimizer_config.optimizer,
# "scheduler": camera_optimizer_config.scheduler,
# }
return Optimizers(optimizer_config, param_groups)
def train(self) -> None:

21
nerfstudio/models/nerfacto.py
View File

@ -67,7 +67,7 @@ class NerfactoModelConfig(ModelConfig):
"""How far along the ray to start sampling."""
far_plane: float = 1000.0
"""How far along the ray to stop sampling."""
background_color: Literal["background", "last_sample"] = "last_sample"
background_color: Literal["random", "last_sample"] = "last_sample"
"""Whether to randomize the background color."""
num_levels: int = 16
"""Number of levels of the hashmap for the base mlp."""
@ -116,6 +116,10 @@ class NerfactoModelConfig(ModelConfig):
"""Whether use single jitter or not for the proposal networks."""
predict_normals: bool = False
"""Whether to predict normals or not."""
use_scene_contraction: bool = True
"""Whether to use scene contraction or not."""
scene_contraction_norm: Literal["inf", "l2"] = "inf"
"""Which norm to use for the scene contraction."""
class NerfactoModel(Model):
@ -131,7 +135,15 @@ class NerfactoModel(Model):
"""Set the fields and modules."""
super().populate_modules()
scene_contraction = SceneContraction(order=float("inf"))
# 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
else:
raise ValueError("Invalid scene contraction norm")
scene_contraction = SceneContraction(order=order)
# Fields
self.field = TCNNNerfactoField(
@ -152,8 +164,9 @@ class NerfactoModel(Model):
# Build the proposal network(s)
self.proposal_networks = torch.nn.ModuleList()
if self.config.use_same_proposal_network:
assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
prop_net_args = self.config.proposal_net_args_list[0]
# assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
# using the last defined proposal network
prop_net_args = self.config.proposal_net_args_list[-1]
network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args)
self.proposal_networks.append(network)
self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])

12
nerfstudio/pipelines/base_pipeline.py
View File

@ -203,7 +203,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
"""Whether to optimize the appearance during evaluation."""
eval_num_appearance_iters: int = 10
"""Number of iterations to optimize the appearance."""
eval_image_scale_factor: float = 0.25
eval_image_scale_factor: float = 0.125
"""Scale factor to use for evaluation images so that they fit in memory."""
@ -367,11 +367,11 @@ class VanillaPipeline(Pipeline):
param_groups = self.datamanager.get_param_groups()
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":
raise ValueError(
"You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses."
)
# # check that camera optimizer is on
# if self.config.datamanager.camera_optimizer.mode == "off":
# raise ValueError(
# "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses."
# )
# CONSOLE.print("Optimizing camera poses...")
# rescale the image to 1/4 resolution

1
pyproject.toml
View File

@ -20,6 +20,7 @@ dependencies = [
"av==9.2.0",
"tyro>=0.3.31",
"gdown==4.5.1",
"GPUtil>=1.4.0",
"ninja==1.10.2.3",
"functorch==0.2.1",
"h5py>=2.9.0",

82
scripts/benchmarking/benchmark_nerfacto.py
View File

@ -4,6 +4,11 @@ Process the Mipnerf360 dataset.
from nerfstudio.utils.scripts import run_command
import GPUtil
import time
import threading
capture_names = [
"Egypt",
"person",
@ -17,36 +22,87 @@ capture_names = [
"Giannini-Hall",
]
table_rows = [
("nerfacto", ""),
("w/o-pose", "--pipeline.datamanager.camera-optimizer.mode off"),
("w/o-app", "--pipeline.model.use-appearance-embedding False"),
("nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True"),
(
"w/o-pose",
"--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.datamanager.camera-optimizer.mode off",
),
(
"w/o-app",
"--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.model.use-appearance-embedding False",
),
(
"w/o-pose-app",
"--pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False",
"--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",
"--pipeline.model.num-proposal-samples-per-ray \"256\" --pipeline.model.num_proposal_iterations 1",
'--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",
"--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.scene-contraction-norm l2",
),
(
"shared-prop-network",
"--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-same-proposal-network True",
),
(
"random-background-color",
"--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.background-color random",
),
]
# 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:
# 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}",
"--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,
)
)
print(command)
jobs.append(command)
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)
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)
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()
# 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()
print("Finished all threads")

36
scripts/benchmarking/create_tables.py Normal file
View File

@ -0,0 +1,36 @@
"""
Code to create the tables.
"""
import pandas as pd
# from nerfstudio.scripts.benchmarking.benchmark_nerfacto import table_rows
# the csv filenames
csv_filename = "wandb-psnr.csv"
# code to read the csv file
df = pd.read_csv(csv_filename)
column_titles = list(df.columns.values)
last_row = list(df.values[-1])
print(len(column_titles))
print(len(last_row))
from collections import defaultdict
results = defaultdict(dict)
for column_title, last_row_value in zip(column_titles, last_row):
if column_title == "Step" or "__MAX" in column_title or "__MIN" in column_title:
continue
name = column_title.split(" - ")[0]
capture_name, method = name.split("_")
metric = float(last_row_value)
# format to have 2 decimal places
metric = f"{metric:.2f}"
results[capture_name][method] = metric
import pprint
pprint.pprint(results)
print(len(results))

19
scripts/downloads/download_data.py
View File

@ -97,6 +97,15 @@ nerfstudio_dataset = [
"sculpture",
"Giannini-Hall",
]
mipnerf360_dataset = [
"mipnerf-360-bicycle",
"mipnerf-360-garden",
"mipnerf-360-stump",
"mipnerf-360-room",
"mipnerf-360-counter",
"mipnerf-360-kitchen",
"mipnerf-360-bonsai",
]
nerfstudio_file_ids = {
"bww_entrance": grab_file_id("https://drive.google.com/file/d/1ylkRHtfB3n3IRLf2wplpfxzPTq7nES9I/view?usp=sharing"),
"campanile": grab_file_id("https://drive.google.com/file/d/13aOfGJRRH05pOOk9ikYGTwqFc2L1xskU/view?usp=sharing"),
@ -106,6 +115,7 @@ nerfstudio_file_ids = {
"redwoods2": grab_file_id("https://drive.google.com/file/d/1rg-4NoXT8p6vkmbWxMOY6PSG4j3rfcJ8/view?usp=sharing"),
"storefront": grab_file_id("https://drive.google.com/file/d/16b792AguPZWDA_YC4igKCwXJqW0Tb21o/view?usp=sharing"),
"vegetation": grab_file_id("https://drive.google.com/file/d/1wBhLQ2odycrtU39y2akVurXEAt9SsVI3/view?usp=sharing"),
# Nerfstudio Dataset
"Egypt": "https://data.nerf.studio/nerfstudio/Egypt.zip",
"person": "https://data.nerf.studio/nerfstudio/person.zip",
"kitchen": "https://data.nerf.studio/nerfstudio/kitchen.zip",
@ -116,8 +126,17 @@ nerfstudio_file_ids = {
"stump": "https://data.nerf.studio/nerfstudio/stump.zip",
"sculpture": "https://data.nerf.studio/nerfstudio/sculpture.zip",
"Giannini-Hall": "https://data.nerf.studio/nerfstudio/Giannini-Hall.zip",
# MipNeRF-360 Dataset. TODO: add these links
"mipnerf360-bicycle": "",
"mipnerf360-garden": "",
"mipnerf360-stump": "",
"mipnerf360-room": "",
"mipnerf360-counter": "",
"mipnerf360-kitchen": "",
"mipnerf360-bonsai": "",
"all": None,
"nerfstudio-dataset": nerfstudio_dataset,
"mipnerf360-dataset": mipnerf360_dataset,
}
NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys())