4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00

adding nerfactory data download option (#565)

* adding nerfactory data download option

* adding docstring

* adding google file id dict
This commit is contained in:
Terrance Wang 2022-09-28 16:10:19 -07:00 committed by GitHub
parent e8a5941b24
commit f0adbba91f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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__":