From f0adbba91fcfa0150ca43e70eefb85313c84ef65 Mon Sep 17 00:00:00 2001
From: Terrance Wang <terrance.3512@gmail.com>
Date: Wed, 28 Sep 2022 16:10:19 -0700
Subject: [PATCH] adding nerfactory data download option (#565)

* adding nerfactory data download option

* adding docstring

* adding google file id dict
---
 scripts/downloads/download_data.py | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py
index 61e643f2b..64758c9e9 100755
--- a/scripts/downloads/download_data.py
+++ b/scripts/downloads/download_data.py
@@ -2,11 +2,18 @@
 """Download datasets"""
 import zipfile
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Optional
 
 import dcargs
 import gdown
 
+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"""
@@ -63,7 +70,20 @@ def download_blender():
     download_path.unlink(missing_ok=True)
 
 
-def main(dataset: Literal["blender", "friends"]):
+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")
+    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)
+
+
+def main(
+    dataset: Literal["blender", "friends", "nerfactory"],
+    dataset_name: Optional[str] = None,
+):
     """Main download script to download all data"""
     save_dir = Path("data/")
     save_dir.mkdir(parents=True, exist_ok=True)
@@ -72,6 +92,12 @@ def main(dataset: Literal["blender", "friends"]):
         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()}"
+            )
+        download_nerfactory(dataset_name)
 
 
 if __name__ == "__main__":