From 1eae0839b44d0b38c8d33c6bac1dafee5bd43d88 Mon Sep 17 00:00:00 2001 From: Jake Austin <55679937+jake-austin@users.noreply.github.com> Date: Fri, 25 Nov 2022 21:40:13 -0800 Subject: [PATCH] Updating the script for downloading data (#968) * Updating the script for downloading data * Making update to config structure * Updating documentation and colab * Adding back the save directory * Pushing an all flag for nerfstudio * forgot to save some previous edits --- README.md | 2 +- colab/demo.ipynb | 50 ++--- docs/nerfology/methods/semantic_nerfw.md | 2 +- docs/quickstart/first_nerf.md | 2 +- scripts/downloads/download_data.py | 264 +++++++++++++++-------- 5 files changed, 208 insertions(+), 112 deletions(-) diff --git a/README.md b/README.md index 114568e10..beda2bb56 100644 --- a/README.md +++ b/README.md @@ -125,7 +125,7 @@ The following will train a _nerfacto_ model, our recommended model for real worl ```bash # Download some test data: -ns-download-data --dataset=nerfstudio --capture=poster +ns-download-data nerfstudio --capture-name=poster # Train model ns-train nerfacto --data data/nerfstudio/poster ``` diff --git a/colab/demo.ipynb b/colab/demo.ipynb index ef1cb4dc2..559299323 100644 --- a/colab/demo.ipynb +++ b/colab/demo.ipynb @@ -3,8 +3,8 @@ { "cell_type": "markdown", "metadata": { - "id": "view-in-github", - "colab_type": "text" + "colab_type": "text", + "id": "view-in-github" }, "source": [ "\"Open" @@ -65,17 +65,17 @@ "cell_type": "code", "execution_count": 1, "metadata": { + "cellView": "form", "colab": { "base_uri": "https://localhost:8080/" }, "id": "RGr33zHaHak0", - "outputId": "2641dce6-cd57-4d31-d8b1-860bc1adf17f", - "cellView": "form" + "outputId": "2641dce6-cd57-4d31-d8b1-860bc1adf17f" }, "outputs": [ { - "output_type": "stream", "name": "stdout", + "output_type": "stream", "text": [ "⏬ Downloading https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh...\n", "📦 Installing...\n", @@ -98,8 +98,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "id": "9oyLHl8QfYwP", - "cellView": "form" + "cellView": "form", + "id": "9oyLHl8QfYwP" }, "outputs": [], "source": [ @@ -127,8 +127,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "id": "msVLprI4gRA4", - "cellView": "form" + "cellView": "form", + "id": "msVLprI4gRA4" }, "outputs": [], "source": [ @@ -194,7 +194,7 @@ " scene = \"custom_data\"\n", "else:\n", " %cd /content/\n", - " !ns-download-data --dataset=nerfstudio --capture=$scene\n", + " !ns-download-data nerfstudio --capture-name=$scene\n", "\n", "print(\"Data Processing Succeeded!\")" ] @@ -203,18 +203,18 @@ "cell_type": "code", "execution_count": 2, "metadata": { + "cellView": "form", "colab": { "base_uri": "https://localhost:8080/", "height": 928 }, "id": "VoKDxqEcjmfC", - "outputId": "d2919aa4-96dd-4e50-829f-289e4208882b", - "cellView": "form" + "outputId": "d2919aa4-96dd-4e50-829f-289e4208882b" }, "outputs": [ { - "output_type": "stream", "name": "stdout", + "output_type": "stream", "text": [ "/content\n", "\u001b[K\u001b[?25h/tools/node/bin/lt -> /tools/node/lib/node_modules/localtunnel/bin/lt.js\n", @@ -225,11 +225,7 @@ ] }, { - "output_type": "execute_result", "data": { - "text/plain": [ - "" - ], "text/html": [ "\n", " \n", " " + ], + "text/plain": [ + "" ] }, + "execution_count": 2, "metadata": {}, - "execution_count": 2 + "output_type": "execute_result" } ], "source": [ @@ -281,8 +281,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "id": "m_N8_cLfjoXD", - "cellView": "form" + "cellView": "form", + "id": "m_N8_cLfjoXD" }, "outputs": [], "source": [ @@ -362,22 +362,22 @@ "accelerator": "GPU", "colab": { "collapsed_sections": [], - "provenance": [], - "include_colab_link": true + "include_colab_link": true, + "provenance": [] }, "gpuClass": "standard", "kernelspec": { - "display_name": "Python 3.8.12 ('nerfactory')", + "display_name": "Python 3.8.13 ('nerfstudio')", "language": "python", "name": "python3" }, "language_info": { "name": "python", - "version": "3.8.12" + "version": "3.8.13" }, "vscode": { "interpreter": { - "hash": "34c28001ff35fb390494047002768a8182dcf55b1b11415165e62ea61557ab83" + "hash": "c59f626636933ef1dc834fb3684b382f705301c5306cf8436d2da634c2289783" } } }, diff --git a/docs/nerfology/methods/semantic_nerfw.md b/docs/nerfology/methods/semantic_nerfw.md index 4037be577..3888fdf8a 100644 --- a/docs/nerfology/methods/semantic_nerfw.md +++ b/docs/nerfology/methods/semantic_nerfw.md @@ -25,7 +25,7 @@ The transient embeddings are still under development. Please stay tuned. Download the Friends Dataset ```bash -ns-download-data --dataset friends +ns-download-data friends ``` ```bash diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index bb9ca8b3e..1bd6ec367 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -6,7 +6,7 @@ The following will train a _nerfacto_ model, our recommended model for real worl ```bash # Download some test data: -ns-download-data --dataset=nerfstudio --capture=poster +ns-download-data nerfstudio --capture-name=poster # Train model ns-train nerfacto --data data/nerfstudio/poster ``` diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 0c37255ab..e9349140e 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -1,54 +1,81 @@ """Download datasets and specific captures from the datasets.""" +from __future__ import annotations + import os import shutil -import sys +import tarfile import zipfile +from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import Union import gdown import tyro from rich.console import Console -from typing_extensions import Literal +from typing_extensions import Annotated + +from nerfstudio.configs.base_config import PrintableConfig CONSOLE = Console(width=120) -def download_blender(save_dir: Path): +@dataclass +class DatasetDownload(PrintableConfig): + """Download a dataset""" + + capture_name = None + + save_dir: Path = Path("data/") + """The directory to save the dataset to""" + + def download(self, save_dir: Path) -> None: + """Download the dataset""" + raise NotImplementedError + + +@dataclass +class BlenderDownload(DatasetDownload): """Download the blender dataset.""" - # TODO: give this code the same structure as download_nerfstudio - # https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG - blender_file_id = "18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG" + def download(self, save_dir: Path): + """Download the blender dataset.""" + # TODO: give this code the same structure as download_nerfstudio - final_path = save_dir / Path("blender") - if os.path.exists(final_path): - shutil.rmtree(str(final_path)) - url = f"https://drive.google.com/uc?id={blender_file_id}" - download_path = save_dir / "blender_data.zip" - gdown.download(url, output=str(download_path)) - with zipfile.ZipFile(download_path, "r") as zip_ref: - zip_ref.extractall(str(save_dir)) - unzip_path = save_dir / Path("nerf_synthetic") - final_path = save_dir / Path("blender") - unzip_path.rename(final_path) - if download_path.exists(): - download_path.unlink() + # https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG + blender_file_id = "18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG" + + final_path = save_dir / Path("blender") + if os.path.exists(final_path): + shutil.rmtree(str(final_path)) + url = f"https://drive.google.com/uc?id={blender_file_id}" + download_path = save_dir / "blender_data.zip" + gdown.download(url, output=str(download_path)) + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(str(save_dir)) + unzip_path = save_dir / Path("nerf_synthetic") + final_path = save_dir / Path("blender") + unzip_path.rename(final_path) + if download_path.exists(): + download_path.unlink() -def download_friends(save_dir: Path): +@dataclass +class FriendsDownload(DatasetDownload): """Download the friends dataset.""" - # https://drive.google.com/file/d/1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5/view?usp=sharing - friends_file_id = "1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5" + def download(self, save_dir: Path): + """Download the friends dataset.""" - # Download the files - url = f"https://drive.google.com/uc?id={friends_file_id}" - download_path = str(save_dir / "friends.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) + # 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 = str(save_dir / "friends.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: @@ -67,9 +94,10 @@ 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"), + "all": None, } -DatasetName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys()) +NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys()) def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str, capture_name_to_file_id: dict): @@ -94,43 +122,132 @@ def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str, os.remove(download_path) -def download_nerfstudio(save_dir: Path, capture_name: str): - """Download specific captures for the nerfstudio dataset.""" - download_capture_name(save_dir, "nerfstudio", capture_name, capture_name_to_file_id=nerfstudio_file_ids) +@dataclass +class NerfstudioDownload(DatasetDownload): + """Download the nerfstudio dataset.""" + + capture_name: NerfstudioCaptureName = "bww_entrance" + + def download(self, save_dir: Path): + """Download the nerfstudio dataset.""" + if self.capture_name == "all": + for capture_name in nerfstudio_file_ids: + if capture_name != "all": + download_capture_name(save_dir, "nerfstudio", capture_name, nerfstudio_file_ids) + return + + download_capture_name(save_dir, "nerfstudio", self.capture_name, capture_name_to_file_id=nerfstudio_file_ids) record3d_file_ids = { "bear": grab_file_id("https://drive.google.com/file/d/1WRZohWMRj0nNlYFIEBwkddDoGPvLTzkR/view?usp=sharing") } - -def download_record3d(save_dir: Path, capture_name: str): - """Download specific captures for the record3d dataset.""" - download_capture_name(save_dir, "record3d", capture_name, capture_name_to_file_id=record3d_file_ids) +Record3dCaptureName = tyro.extras.literal_type_from_choices(record3d_file_ids.keys()) -def download_dnerf(save_dir: Path): - """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF).""" - # TODO: give this code the same structure as download_nerfstudio +@dataclass +class Record3dDownload(DatasetDownload): + """Download the record3d dataset.""" - final_path = save_dir / Path("dnerf") - if os.path.exists(final_path): - shutil.rmtree(str(final_path)) - download_path = save_dir / "dnerf_data.zip" - os.system(f"curl -L https://www.dropbox.com/s/raw/0bf6fl0ye2vz3vr/data.zip > {download_path}") - with zipfile.ZipFile(download_path, "r") as zip_ref: - zip_ref.extractall(str(save_dir)) - unzip_path = save_dir / Path("data") - final_path = save_dir / Path("dnerf") - unzip_path.rename(final_path) - if download_path.exists(): - download_path.unlink() + capture_name: Record3dCaptureName = "bear" + + def download(self, save_dir: Path): + download_capture_name(save_dir, "record3d", self.capture_name, capture_name_to_file_id=record3d_file_ids) + + +@dataclass +class DNerfDownload(DatasetDownload): + """Download the dnerf dataset.""" + + def download(self, save_dir: Path): + """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF).""" + # TODO: give this code the same structure as download_nerfstudio + + final_path = save_dir / Path("dnerf") + if os.path.exists(final_path): + shutil.rmtree(str(final_path)) + download_path = save_dir / "dnerf_data.zip" + os.system(f"curl -L https://www.dropbox.com/s/raw/0bf6fl0ye2vz3vr/data.zip > {download_path}") + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(str(save_dir)) + unzip_path = save_dir / Path("data") + final_path = save_dir / Path("dnerf") + unzip_path.rename(final_path) + if download_path.exists(): + download_path.unlink() + + +# pylint: disable=line-too-long +phototourism_downloads = { + "brandenburg-gate": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/brandenburg_gate.tar.gz", + "buckingham-palace": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/buckingham_palace.tar.gz", + "colosseum-exterior": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/colosseum_exterior.tar.gz", + "grand-palace-brussels": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/grand_place_brussels.tar.gz", + "notre-dame-facade": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/notre_dame_front_facade.tar.gz", + "westminster-palace": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/palace_of_westminster.tar.gz", + "pantheon-exterior": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/pantheon_exterior.tar.gz", + "taj-mahal": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/taj_mahal.tar.gz", + "temple-nara": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/temple_nara_japan.tar.gz", + "trevi-fountain": "https://www.cs.ubc.ca/research/kmyi_data/imw2020/TrainingData/trevi_fountain.tar.gz", + "all": None, +} + +PhototourismCaptureName = tyro.extras.literal_type_from_choices(phototourism_downloads.keys()) + + +@dataclass +class PhototourismDownload(DatasetDownload): + """Download the phototourism dataset.""" + + capture_name: PhototourismCaptureName = "brandenburg-gate" + + def download(self, save_dir: Path): + """Download a PhotoTourism dataset: https://www.cs.ubc.ca/~kmyi/imw2020/data.html""" + + if self.capture_name == "all": + for capture_name in phototourism_downloads: + if capture_name != "all": + PhototourismDownload(capture_name=capture_name).download(save_dir) + return + + assert ( + self.capture_name in phototourism_downloads + ), f"Capture name {self.capture_name} not found in {phototourism_downloads.keys()}" + url = phototourism_downloads[self.capture_name] + target_path = str(save_dir / f"phototourism/{self.capture_name}") + os.makedirs(target_path, exist_ok=True) + download_path = Path(f"{target_path}.tar.gz") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + + os.system(f"curl -L {url} > {download_path}") + + with tarfile.open(download_path, "r:gz") as tar_ref: + tar_ref.extractall(str(tmp_path)) + + inner_folders = os.listdir(tmp_path) + assert len(inner_folders) == 1, "There is more than one folder inside this zip file." + folder = os.path.join(tmp_path, inner_folders[0]) + shutil.rmtree(target_path) + shutil.move(folder, target_path) + shutil.rmtree(tmp_path) + os.remove(download_path) + + +Commands = Union[ + Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], + Annotated[FriendsDownload, tyro.conf.subcommand(name="friends")], + Annotated[NerfstudioDownload, tyro.conf.subcommand(name="nerfstudio")], + Annotated[Record3dDownload, tyro.conf.subcommand(name="record3d")], + Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")], + Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")], +] def main( - dataset: Literal["blender", "friends", "nerfstudio", "record3d", "dnerf"], - capture_name: Optional[DatasetName] = None, # type: ignore - save_dir: Path = Path("data/"), + dataset: DatasetDownload, ): """Script to download existing datasets. @@ -140,47 +257,26 @@ def main( which capture to download. - blender: Blender synthetic scenes realeased with NeRF. - friends: Friends TV show scenes. + - record3d: Record3d dataset. + - dnerf: D-NeRF dataset. + - phototourism: PhotoTourism dataset. Use the `capture_name` argument to specify which capture to download. Args: dataset: The dataset to download (from). - capture_name: The capture name to download (if downloading from nerfstudio dataset). - save_dir: The directory to save the data to. """ - save_dir.mkdir(parents=True, exist_ok=True) + dataset.save_dir.mkdir(parents=True, exist_ok=True) - if dataset == "blender": - if capture_name is not None: - CONSOLE.print("Capture name is ignored when downloading from the blender dataset.") - download_blender(save_dir) - if dataset == "friends": - if capture_name is not None: - CONSOLE.print("Capture name is ignored when downloading from the blender dataset.") - download_friends(save_dir) - if dataset == "nerfstudio": - if capture_name is None: - capture_names = sorted(nerfstudio_file_ids.keys()) - CONSOLE.rule("[bold red]Error", style="bold red") - CONSOLE.print("[bold yellow]You must pass in --capture-name when downloading from the nerfstudio dataset.") - CONSOLE.print("Use one of the following:") - CONSOLE.print(f"\t {capture_names}") - sys.exit() - download_nerfstudio(save_dir, capture_name) - if dataset == "record3d": - download_record3d(save_dir, "bear") - if dataset == "dnerf": - if capture_name is not None: - CONSOLE.print("Capture name is ignored when downloading from the dnerf dataset.") - download_dnerf(save_dir) + dataset.download(dataset.save_dir) def entrypoint(): """Entrypoint for use with pyproject scripts.""" tyro.extras.set_accent_color("bright_yellow") - tyro.cli(main) + main(tyro.cli(Commands)) if __name__ == "__main__": entrypoint() # For sphinx docs -get_parser_fn = lambda: tyro.extras.get_parser(main) # noqa +get_parser_fn = lambda: tyro.extras.get_parser(Commands) # noqa