From ce8985aeb62c0ccdfa9698699e01fb4c6fb1c464 Mon Sep 17 00:00:00 2001
From: Matthew Tancik <tancik@berkeley.edu>
Date: Tue, 4 Oct 2022 21:00:21 -0700
Subject: [PATCH] cli improvements to download script (#668)

---
 docs/reference/cli/index.rst           |  1 +
 docs/reference/cli/ns_download_data.md |  9 +++++++++
 pyproject.toml                         |  2 +-
 scripts/downloads/download_data.py     | 27 +++++++++++++++-----------
 4 files changed, 27 insertions(+), 12 deletions(-)
 create mode 100644 docs/reference/cli/ns_download_data.md
 mode change 100755 => 100644 scripts/downloads/download_data.py

diff --git a/docs/reference/cli/index.rst b/docs/reference/cli/index.rst
index f99746fd6..a473247ab 100644
--- a/docs/reference/cli/index.rst
+++ b/docs/reference/cli/index.rst
@@ -7,5 +7,6 @@ CLI
     :maxdepth: 1
 
     ns_process_data
+    ns_download_data
     ns_render
     ns_eval
\ No newline at end of file
diff --git a/docs/reference/cli/ns_download_data.md b/docs/reference/cli/ns_download_data.md
new file mode 100644
index 000000000..60b4595f3
--- /dev/null
+++ b/docs/reference/cli/ns_download_data.md
@@ -0,0 +1,9 @@
+# ns-download-data
+
+```{eval-rst}
+.. argparse::
+    :module: scripts.downloads.download_data
+    :func: get_parser_fn
+    :prog: ns-process-data
+    :nodefault:
+```
diff --git a/pyproject.toml b/pyproject.toml
index 3582410cb..a5030b879 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,7 @@ dependencies = [
     "aiortc==1.3.2",
     "appdirs>=1.4",
     "av==9.2.0",
-    "dcargs>=0.3.17",
+    "dcargs>=0.3.20",
     "gdown==4.5.1",
     "ninja==1.10.2.3",
     "functorch==0.2.1",
diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py
old mode 100755
new mode 100644
index e5255d6f2..f6a42b412
--- a/scripts/downloads/download_data.py
+++ b/scripts/downloads/download_data.py
@@ -60,7 +60,7 @@ nerfstudio_file_ids = {
     "bunny": grab_file_id("https://drive.google.com/file/d/1oVytSgd1gQclF0CybJ6JdXDmaqKOx7Wg/view?usp=sharing"),
     "bww_entrance": grab_file_id("https://drive.google.com/file/d/1ylkRHtfB3n3IRLf2wplpfxzPTq7nES9I/view?usp=sharing"),
     "bww_tree": grab_file_id("https://drive.google.com/file/d/1N_OQejT0MblK1UP05KORxy25SXjKqZN1/view?usp=sharing"),
-    "campanelle": grab_file_id("https://drive.google.com/file/d/13aOfGJRRH05pOOk9ikYGTwqFc2L1xskU/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"),
     "japanese_maple": grab_file_id(
@@ -75,6 +75,8 @@ nerfstudio_file_ids = {
     "vegetation": grab_file_id("https://drive.google.com/file/d/1wBhLQ2odycrtU39y2akVurXEAt9SsVI3/view?usp=sharing"),
 }
 
+DatasetName = dcargs.extras.literal_type_from_choices(nerfstudio_file_ids.keys())
+
 
 def download_nerfstudio(save_dir: Path, capture_name: str):
     """Download specific captures from the nerfstudio dataset."""
@@ -100,33 +102,33 @@ def download_nerfstudio(save_dir: Path, capture_name: str):
 
 def main(
     dataset: Literal["blender", "friends", "nerfstudio"],
-    capture_name: Optional[str] = None,
+    capture_name: Optional[DatasetName] = None,  # type: ignore
     save_dir: Path = Path("data/"),
 ):
     """Main download script to download all data.
 
     Args:
         dataset: The dataset to download (from).
-        capture_name: The capture name to download (from the dataset).
+        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)
 
     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.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}"
-            )
-            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}")
+            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)
 
@@ -139,3 +141,6 @@ def entrypoint():
 
 if __name__ == "__main__":
     entrypoint()
+
+# For sphinx docs
+get_parser_fn = lambda: dcargs.extras.get_parser(main)  # noqa