diff --git a/.gitignore b/.gitignore index 64742e3ce..51c91f933 100644 --- a/.gitignore +++ b/.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/ diff --git a/nerfstudio/datamanagers/dataparsers/nerfactory_parser.py b/nerfstudio/datamanagers/dataparsers/nerfactory_parser.py index 75d2111a6..af276954c 100644 --- a/nerfstudio/datamanagers/dataparsers/nerfactory_parser.py +++ b/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.""" diff --git a/nerfstudio/optimizers/optimizers.py b/nerfstudio/optimizers/optimizers.py index eee89c74f..744a698d8 100644 --- a/nerfstudio/optimizers/optimizers.py +++ b/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 = {} diff --git a/nerfstudio/utils/colors.py b/nerfstudio/utils/colors.py index 9b5688ccf..938031f00 100644 --- a/nerfstudio/utils/colors.py +++ b/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) diff --git a/scripts/docs/add_nb_tags.py b/scripts/docs/add_nb_tags.py index ea4468c8f..f967a35ba 100755 --- a/scripts/docs/add_nb_tags.py +++ b/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 diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index a4dafaf9f..ca4bc697a 100755 --- a/scripts/downloads/download_data.py +++ b/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(): diff --git a/scripts/process_data.py b/scripts/process_data.py index d700d78d4..23a8160a3 100755 --- a/scripts/process_data.py +++ b/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(