don't download to /tmp (#612)

scripts/downloads/download_data.py

@@ -13,24 +13,27 @@ from rich.console import Console
 console = Console(width=120)
 
 
-def download_blender():
+def download_blender(save_dir: Path):
     """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"
 
+    final_path = save_dir / Path("blender")
+    os.remove(str(final_path))
     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)
+    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("data/")
-    unzip_path = Path("data/nerf_synthetic")
-    final_path = Path("data/blender")
+        zip_ref.extractall(str(save_dir))
+    unzip_path = save_dir / Path("nerf_synthetic")
+    final_path = save_dir / Path("blender")
     unzip_path.rename(final_path)
     download_path.unlink(missing_ok=True)
 
 
-def download_friends():
+def download_friends(save_dir: Path):
     """Download the friends dataset."""
 
     # https://drive.google.com/file/d/1sgKr0ZO7BQC0FYinAnRSxobIWNucAST5/view?usp=sharing
@@ -38,10 +41,10 @@ def download_friends():
 
     # 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)
+    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("data/")
+        zip_ref.extractall(str(save_dir))
     os.remove(download_path)
 
 
@@ -53,37 +56,46 @@ nerfstudio_file_ids = {
 }
 
 
-def download_nerfstudio(capture_name: str):
+def download_nerfstudio(save_dir: Path, 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}"
+    target_path = str(save_dir / f"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)
+    tmp_path = str(save_dir / ".temp")
+    shutil.rmtree(tmp_path, ignore_errors=True)
+    os.makedirs(tmp_path, exist_ok=True)
+    gdown.download(url, output=str(download_path))
     with zipfile.ZipFile(download_path, "r") as zip_ref:
-        zip_ref.extractall(f"/tmp/{target_path}")
-    inner_folders = os.listdir(f"/tmp/{target_path}")
+        zip_ref.extractall(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(f"/tmp/{target_path}", inner_folders[0])
+    folder = os.path.join(tmp_path, inner_folders[0])
     shutil.rmtree(target_path)
     shutil.move(folder, target_path)
-    shutil.rmtree(f"/tmp/{target_path}")
+    shutil.rmtree(tmp_path)
     os.remove(download_path)
 
 
 def main(
     dataset: Literal["blender", "friends", "nerfstudio"],
     capture_name: Optional[str] = None,
+    save_dir: Path = Path("data/"),
 ):
-    """Main download script to download all data"""
-    save_dir = 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).
+        save_dir: The directory to save the data to.
+    """
     save_dir.mkdir(parents=True, exist_ok=True)
 
     if dataset == "blender":
-        download_blender()
+        download_blender(save_dir)
     if dataset == "friends":
-        download_friends()
+        download_friends(save_dir)
     if dataset == "nerfstudio":
         if capture_name is None:
             capture_names = sorted(nerfstudio_file_ids.keys())
@@ -96,7 +108,7 @@ def main(
             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)
+        download_nerfstudio(save_dir, capture_name)
 
 
 def entrypoint():