updated download script (#592)

* updated download script

* rename

* remove omegaconf

.gitignore

@@ -168,12 +168,10 @@ outputs/
 events.out.*
 
 # Data
-data
+data/
 data_*/
 !nerfstudio/data
-!configs/data
-configs_custom
-!docs/tutorials/data
+!docs/tutorials/data/
 
 # Misc
 old/

nerfstudio/datamanagers/dataparsers/nerfactory_parser.py

@@ -37,7 +37,7 @@ class NerfactoryDataParserConfig(DataParserConfig):
 
     _target: Type = field(default_factory=lambda: Nerfactory)
     """target class to instantiate"""
-    data_directory: Path = Path("data/ours/posters_v3")
+    data_directory: Path = Path("data/nerfstudio/posters_v3")
     """directory specifying location of data"""
     scale_factor: float = 1.0
     """How much to scale the camera origins by."""

nerfstudio/optimizers/optimizers.py

@@ -19,16 +19,13 @@ from __future__ import annotations
 
 from typing import Any, Dict, List, Union
 
-from omegaconf import DictConfig
 from torch.cuda.amp.grad_scaler import GradScaler
 from torch.nn.parameter import Parameter
 
 from nerfstudio.utils import writer
 
 
-def setup_optimizers(
-    config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]]
-) -> "Optimizers":
+def setup_optimizers(config: Union[Dict[str, Any]], param_groups: Dict[str, List[Parameter]]) -> "Optimizers":
     """Helper to set up the optimizers
 
     Args:
@@ -49,7 +46,7 @@ class Optimizers:
         param_groups: A dictionary of parameter groups to optimize.
     """
 
-    def __init__(self, config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]]):
+    def __init__(self, config: Union[Dict[str, Any]], param_groups: Dict[str, List[Parameter]]):
         self.config = config
         self.optimizers = {}
         self.schedulers = {}

nerfstudio/utils/colors.py

@@ -16,7 +16,6 @@
 from typing import Union
 
 import torch
-from omegaconf import ListConfig
 from torchtyping import TensorType
 
 WHITE = torch.tensor([1.0, 1.0, 1.0])
@@ -34,10 +33,10 @@ COLORS_DICT = {
 }
 
 
-def get_color(color: Union[str, list, ListConfig]) -> TensorType[3]:
+def get_color(color: Union[str, list]) -> TensorType[3]:
     """
     Args:
-        color (Union[str, list, ListConfig]): Color as a string or a rgb list
+        color (Union[str, list]): Color as a string or a rgb list
 
     Returns:
         TensorType[3]: Parsed color
@@ -47,7 +46,7 @@ def get_color(color: Union[str, list, ListConfig]) -> TensorType[3]:
         if color not in COLORS_DICT:
             raise ValueError(f"{color} is not a valid preset color")
         return COLORS_DICT[color]
-    if isinstance(color, (list, ListConfig)):
+    if isinstance(color, list):
         if len(color) != 3:
             raise ValueError(f"Color should be 3 values (RGB) instead got {color}")
         return torch.tensor(color)

scripts/docs/add_nb_tags.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python
 """Helper that add tags to notebooks based on cell comments."""
 
 import sys

scripts/downloads/download_data.py

@@ -1,65 +1,25 @@
-#!/usr/bin/env python
-"""Download datasets"""
+"""Download datasets and specific captures from the datasets."""
+import os
+import shutil
+import sys
 import zipfile
 from pathlib import Path
 from typing import Literal, Optional
 
 import dcargs
 import gdown
+from rich.console import Console
 
-nerfactory_file_ids = {
-    "dozer": "1-OR5F_V5S4s-yzxohbwTylaXjzYLu8ZR",
-    "sf_street": "1DbLyptL6my2QprEVtYuW2uzgp9JAK5Wz",
-    "poster": "1dmjWGXlJnUxwosN6MVooCDQe970PkD-1",
-    "lion": "1bx0aDtHoqB3xTAVHL9UIrlTL_vqC6vaH",
-}
-
-
-def download_friends():
-    """Script ported from: https://github.com/ethanweber/sitcoms3D/blob/master/download_data.py"""
-    data_filename_to_file_id = [
-        ("./data/human_data.zip", "1zk12qxnZbReKqdvy-JYp1jQn2QA3oJzz"),
-        ("./data/human_pairs.zip", "1927rnmJ3mWcjsj0afqawSto2IK2lLwmA"),
-        ("./data/sparse_reconstruction_and_nerf_data.zip", "1RmwDUAp1T4RkwZg1S_7L9JAViWtBZfAG"),
-    ]
-
-    disk_data_filename_to_file_id = [  # pylint: disable=unused-variable
-        ("./data/sparse_reconstruction_and_nerf_data/ELR-apartment-disk.zip", "1YcTwMv5PP0uqXYdN-Lp8k0WnXUdHBXrY"),
-        ("./data/sparse_reconstruction_and_nerf_data/Frasier-apartment-disk.zip", "1vPqplDv5rrFu5LrxTHbET0FiePTGi8zE"),
-        (
-            "./data/sparse_reconstruction_and_nerf_data/Friends-monica_apartment-disk.zip",
-            "1yNpU4M44gIuWEaC5tSAOKy-A6KX8wTm8",
-        ),
-        (
-            "./data/sparse_reconstruction_and_nerf_data/HIMYM-red_apartment-disk.zip",
-            "1m76i46YYlpDKvlx-Jn7qDDAvkOVVVmqF",
-        ),
-        (
-            "./data/sparse_reconstruction_and_nerf_data/Seinfeld-jerry_living_room-disk.zip",
-            "1YFoV3Gd9asKsRarZvzRTYtri1755rBPx",
-        ),
-        ("./data/sparse_reconstruction_and_nerf_data/TAAHM-kitchen-disk.zip", "1QpF8kVGfgkqm_cDY5sY_E82BT8VOWL38"),
-        (
-            "./data/sparse_reconstruction_and_nerf_data/TBBT-big_living_room-disk.zip",
-            "175BjkpxMAOt75ZVIjYpArxWpO8uRexgJ",
-        ),
-    ]
-
-    # !!! Large files !!! Only uncomment if you need the disk correspondences.
-    # data_filename_to_file_id += disk_data_filename_to_file_id
-
-    for data_filename, file_id in data_filename_to_file_id:
-        # Download the files
-        url = f"https://drive.google.com/uc?id={file_id}"
-        gdown.download(url, output=data_filename, quiet=False)
-        with zipfile.ZipFile(data_filename, "r") as zip_ref:
-            zip_ref.extractall(data_filename.replace(".zip", "").replace("-disk", ""))
-        Path(data_filename).unlink(missing_ok=True)
+console = Console(width=120)
 
 
 def download_blender():
-    """Download blender dataset and format accordingly"""
-    url = "https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG"
+    """Download the blender dataset."""
+
+    # https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG
+    blender_file_id = "18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG"
+
+    url = f"https://drive.google.com/uc?id={blender_file_id}"
     download_path = Path("data/blender_data.zip")
     gdown.download(url, output=str(download_path), quiet=False)
     with zipfile.ZipFile(download_path, "r") as zip_ref:
@@ -70,19 +30,51 @@ def download_blender():
     download_path.unlink(missing_ok=True)
 
 
-def download_nerfactory(dataset_name: str):
-    """Download a zipped nerfactory dataset"""
-    url = f"https://drive.google.com/uc?id={nerfactory_file_ids[dataset_name]}"
-    download_path = Path(f"data/{dataset_name}.zip")
+def download_friends():
+    """Download the friends dataset."""
+
+    # https://drive.google.com/file/d/1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5/view?usp=sharing
+    friends_file_id = "1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5"
+
+    # Download the files
+    url = f"https://drive.google.com/uc?id={friends_file_id}"
+    download_path = "data/friends.zip"
+    gdown.download(url, output=download_path, quiet=False)
+    with zipfile.ZipFile(download_path, "r") as zip_ref:
+        zip_ref.extractall("data/")
+    os.remove(download_path)
+
+
+# https://drive.google.com/drive/folders/1Wh66z3qQTZ8o2MwPPwYOrdwQtXUEQFyq?usp=sharing
+nerfstudio_file_ids = {
+    "dozer": "1-OR5F_V5S4s-yzxohbwTylaXjzYLu8ZR",
+    "sf_street": "1DbLyptL6my2QprEVtYuW2uzgp9JAK5Wz",
+    "poster": "1dmjWGXlJnUxwosN6MVooCDQe970PkD-1",
+}
+
+
+def download_nerfstudio(capture_name: str):
+    """Download specific captures from the nerfstudio dataset."""
+
+    url = f"https://drive.google.com/uc?id={nerfstudio_file_ids[capture_name]}"
+    target_path = f"data/nerfstudio/{capture_name}"
+    os.makedirs(target_path, exist_ok=True)
+    download_path = Path(f"{target_path}.zip")
     gdown.download(url, output=str(download_path), quiet=False)
     with zipfile.ZipFile(download_path, "r") as zip_ref:
-        zip_ref.extractall(f"data/ours/{dataset_name}")
-    download_path.unlink(missing_ok=True)
+        zip_ref.extractall(f"/tmp/{target_path}")
+    inner_folders = os.listdir(f"/tmp/{target_path}")
+    assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
+    folder = os.path.join(f"/tmp/{target_path}", inner_folders[0])
+    shutil.rmtree(target_path)
+    shutil.move(folder, target_path)
+    shutil.rmtree(f"/tmp/{target_path}")
+    os.remove(download_path)
 
 
 def main(
-    dataset: Literal["blender", "friends", "nerfactory"],
-    dataset_name: Optional[str] = None,
+    dataset: Literal["blender", "friends", "nerfstudio"],
+    capture_name: Optional[str] = None,
 ):
     """Main download script to download all data"""
     save_dir = Path("data/")
@@ -92,12 +84,19 @@ def main(
         download_blender()
     if dataset == "friends":
         download_friends()
-    if dataset == "nerfactory":
-        if dataset_name is None or dataset_name not in nerfactory_file_ids:
-            raise ValueError(
-                f"must pass in dataset_name when downloading nerfactory data, options: {nerfactory_file_ids.keys()}"
+    if dataset == "nerfstudio":
+        if capture_name is None:
+            capture_names = sorted(nerfstudio_file_ids.keys())
+            console.print(
+                "[bold yellow]You must pass in --capture-name when downloading from the nerfstudio dataset."
+                f" Use one of the following: \n\t {capture_names}"
             )
-        download_nerfactory(dataset_name)
+            sys.exit()
+        if capture_name not in nerfstudio_file_ids:
+            capture_names = sorted(nerfstudio_file_ids.keys())
+            console.print(f"[bold yellow]Invalid --capture-name choice. Use one of the following: \n {capture_names}")
+            sys.exit()
+        download_nerfstudio(capture_name)
 
 
 def entrypoint():

scripts/process_data.py

@@ -415,6 +415,7 @@ def main(
     summary_log = []
 
     if data.is_file():
+
         if not verbose:
             with CONSOLE.status("[bold yellow]Converting video to images...", spinner="bouncingBall"):
                 num_vid_frames, num_frames = convert_video_to_images(