merged main

scripts/benchmarking/launch_train_blender.sh

@@ -62,6 +62,13 @@ timestamp=$(date "+%Y-%m-%d_%H%M%S")
 # kill all the background jobs if terminated:
 trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
 
+dataparser="blender-data"
+trans_file=""
+if [ "$method_name" = "instant-ngp-bounded" ]; then
+    dataparser=""
+    trans_file="/transforms_train.json"
+fi
+
 for dataset in "${DATASETS[@]}"; do
     if "$single" && [ -n "${GPU_PID[$idx]+x}" ]; then
         echo "Waiting for GPU ${GPU_IDX[$idx]}"
@@ -70,7 +77,7 @@ for dataset in "${DATASETS[@]}"; do
     fi
     export CUDA_VISIBLE_DEVICES="${GPU_IDX[$idx]}"
     ns-train "${method_name}" "${method_opts[@]}" \
-             --data="data/blender/${dataset}" \
+             --data="data/blender/${dataset}${trans_file}" \
              --experiment-name="blender_${dataset}_${tag}" \
              --trainer.relative-model-dir=nerfstudio_models/ \
              --trainer.steps-per-save=1000 \
@@ -79,7 +86,7 @@ for dataset in "${DATASETS[@]}"; do
              --logging.enable-profiler=False \
              --vis "${vis}" \
              --timestamp "$timestamp" \
-             blender-data & GPU_PID[$idx]=$!
+             ${dataparser} & GPU_PID[$idx]=$!
     echo "Launched ${method_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
     
     # update gpu

scripts/downloads/download_data.py

@@ -15,6 +15,7 @@ from rich.console import Console
 from typing_extensions import Annotated
 
 from nerfstudio.configs.base_config import PrintableConfig
+from nerfstudio.utils.scripts import run_command
 
 CONSOLE = Console(width=120)
 
@@ -84,17 +85,39 @@ def grab_file_id(zip_url: str) -> str:
     return s.split("/")[0]
 
 
+nerfstudio_dataset = [
+    "Egypt",
+    "person",
+    "kitchen",
+    "plane",
+    "dozer",
+    "floating-tree",
+    "aspen",
+    "stump",
+    "sculpture",
+    "Giannini-Hall",
+]
 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"),
     "desolation": grab_file_id("https://drive.google.com/file/d/14IzOOQm9KBJ3kPbunQbUTHPnXnmZus-f/view?usp=sharing"),
-    "dozer": grab_file_id("https://drive.google.com/file/d/1-OR5F_V5S4s-yzxohbwTylaXjzYLu8ZR/view?usp=sharing"),
     "library": grab_file_id("https://drive.google.com/file/d/1Hjbh_-BuaWETQExn2x2qGD74UwrFugHx/view?usp=sharing"),
     "poster": grab_file_id("https://drive.google.com/file/d/1dmjWGXlJnUxwosN6MVooCDQe970PkD-1/view?usp=sharing"),
     "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"),
+    "Egypt": "https://data.nerf.studio/nerfstudio/Egypt.zip",
+    "person": "https://data.nerf.studio/nerfstudio/person.zip",
+    "kitchen": "https://data.nerf.studio/nerfstudio/kitchen.zip",
+    "plane": "https://data.nerf.studio/nerfstudio/plane.zip",
+    "dozer": "https://data.nerf.studio/nerfstudio/dozer.zip",
+    "floating-tree": "https://data.nerf.studio/nerfstudio/floating-tree.zip",
+    "aspen": "https://data.nerf.studio/nerfstudio/aspen.zip",
+    "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",
     "all": None,
+    "nerfstudio-dataset": nerfstudio_dataset,
 }
 
 NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys())
@@ -103,14 +126,33 @@ NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_id
 def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str, capture_name_to_file_id: dict):
     """Download specific captures a given dataset and capture name."""
 
-    url = f"https://drive.google.com/uc?id={capture_name_to_file_id[capture_name]}"
-    target_path = str(save_dir / f"{dataset_name}/{capture_name}")
-    os.makedirs(target_path, exist_ok=True)
-    download_path = Path(f"{target_path}.zip")
-    tmp_path = str(save_dir / ".temp")
-    shutil.rmtree(tmp_path, ignore_errors=True)
-    os.makedirs(tmp_path, exist_ok=True)
-    gdown.download(url, output=str(download_path))
+    file_id_or_zip_url = capture_name_to_file_id[capture_name]
+    if file_id_or_zip_url.endswith(".zip"):
+        url = file_id_or_zip_url  # zip url
+        target_path = str(save_dir / f"{dataset_name}/{capture_name}")
+        os.makedirs(target_path, exist_ok=True)
+        download_path = Path(f"{target_path}.zip")
+        tmp_path = str(save_dir / ".temp")
+        shutil.rmtree(tmp_path, ignore_errors=True)
+        os.makedirs(tmp_path, exist_ok=True)
+        try:
+            os.remove(download_path)
+        except OSError:
+            pass
+        run_command(f"wget {url} -O {download_path}", verbose=True)
+    else:
+        url = f"https://drive.google.com/uc?id={file_id_or_zip_url}"  # file id
+        target_path = str(save_dir / f"{dataset_name}/{capture_name}")
+        os.makedirs(target_path, exist_ok=True)
+        download_path = Path(f"{target_path}.zip")
+        tmp_path = str(save_dir / ".temp")
+        shutil.rmtree(tmp_path, ignore_errors=True)
+        os.makedirs(tmp_path, exist_ok=True)
+        try:
+            os.remove(download_path)
+        except OSError:
+            pass
+        gdown.download(url, output=str(download_path))
     with zipfile.ZipFile(download_path, "r") as zip_ref:
         zip_ref.extractall(tmp_path)
     inner_folders = os.listdir(tmp_path)
@@ -132,7 +174,13 @@ class NerfstudioDownload(DatasetDownload):
         """Download the nerfstudio dataset."""
         if self.capture_name == "all":
             for capture_name in nerfstudio_file_ids:
-                if capture_name != "all":
+                if capture_name not in ("all", "nerfstudio-dataset"):
+                    download_capture_name(save_dir, "nerfstudio", capture_name, nerfstudio_file_ids)
+            return
+
+        if self.capture_name == "nerfstudio-dataset":
+            for capture_name in nerfstudio_dataset:
+                if capture_name not in ("all", "nerfstudio-dataset"):
                     download_capture_name(save_dir, "nerfstudio", capture_name, nerfstudio_file_ids)
             return
 

scripts/render.py

@@ -26,7 +26,6 @@ from typing_extensions import Literal, assert_never
 
 from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path
 from nerfstudio.cameras.cameras import Cameras
-from nerfstudio.configs.base_config import Config  # pylint: disable=unused-import
 from nerfstudio.pipelines.base_pipeline import Pipeline
 from nerfstudio.utils import install_checks
 from nerfstudio.utils.eval_utils import eval_setup

scripts/train.py

@@ -45,8 +45,8 @@ import tyro
 import yaml
 from rich.console import Console
 
-from nerfstudio.configs import base_config as cfg
 from nerfstudio.configs.config_utils import convert_markup_to_ansi
+from nerfstudio.configs.experiment_config import ExperimentConfig
 from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion
 from nerfstudio.engine.trainer import Trainer
 from nerfstudio.utils import comms, profiler
@@ -74,7 +74,7 @@ def _set_random_seed(seed) -> None:
     torch.manual_seed(seed)
 
 
-def train_loop(local_rank: int, world_size: int, config: cfg.Config, global_rank: int = 0):
+def train_loop(local_rank: int, world_size: int, config: ExperimentConfig, global_rank: int = 0):
     """Main training function that sets up and runs the trainer per process
 
     Args:
@@ -95,7 +95,7 @@ def _distributed_worker(
     num_gpus_per_machine: int,
     machine_rank: int,
     dist_url: str,
-    config: cfg.Config,
+    config: ExperimentConfig,
     timeout: timedelta = DEFAULT_TIMEOUT,
 ) -> Any:
     """Spawned distributed worker that handles the initialization of process group and handles the
@@ -110,7 +110,7 @@ def _distributed_worker(
         dist_url: URL to connect to for distributed jobs, including protocol
             E.g., "tcp://127.0.0.1:8686".
             It can be set to "auto" to automatically select a free port on localhost.
-        config: Config specifying training regimen.
+        config: ExperimentConfig specifying training regimen.
         timeout: Timeout of the distributed workers.
 
     Raises:
@@ -150,7 +150,7 @@ def launch(
     num_machines: int = 1,
     machine_rank: int = 0,
     dist_url: str = "auto",
-    config: Optional[cfg.Config] = None,
+    config: Optional[ExperimentConfig] = None,
     timeout: timedelta = DEFAULT_TIMEOUT,
 ) -> None:
     """Function that spawns muliple processes to call on main_func
@@ -161,7 +161,7 @@ def launch(
         num_machines (int, optional): total number of machines
         machine_rank (int, optional): rank of this machine.
         dist_url (str, optional): url to connect to for distributed jobs.
-        config (Config, optional): config file specifying training regimen.
+        config (ExperimentConfig, optional): config file specifying training regimen.
         timeout (timedelta, optional): timeout of the distributed workers.
     """
     assert config is not None
@@ -215,7 +215,7 @@ def launch(
             profiler.flush_profiler(config.logging)
 
 
-def main(config: cfg.Config) -> None:
+def main(config: ExperimentConfig) -> None:
     """Main function."""
 
     config.set_timestamp()