diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index ca4bc697a..dc091ceba 100755 --- a/scripts/downloads/download_data.py +++ b/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():