4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
nerfstudio/scripts/downloads/download_data.py
Brent Yi acd126c1ef
Set dcargs accent color (#465)
* Set dcargs accent color

* Bump required dcargs version

* color -> colour

* Bump dcargs to 0.3.6
2022-09-19 00:28:06 -07:00

85 lines
3.1 KiB
Python
Executable File

#!/usr/bin/env python
"""Download datasets"""
import zipfile
from pathlib import Path
from typing import Literal
import dcargs
import gdown
def download_friends():
"""Script ported from: https://github.com/ethanweber/sitcoms3D/blob/master/download_data.py"""
data_filename_to_file_id = [
("./data/human_data.zip", "1zk12qxnZbReKqdvy-JYp1jQn2QA3oJzz"),
("./data/human_pairs.zip", "1927rnmJ3mWcjsj0afqawSto2IK2lLwmA"),
("./data/sparse_reconstruction_and_nerf_data.zip", "1RmwDUAp1T4RkwZg1S_7L9JAViWtBZfAG"),
]
disk_data_filename_to_file_id = [ # pylint: disable=unused-variable
("./data/sparse_reconstruction_and_nerf_data/ELR-apartment-disk.zip", "1YcTwMv5PP0uqXYdN-Lp8k0WnXUdHBXrY"),
("./data/sparse_reconstruction_and_nerf_data/Frasier-apartment-disk.zip", "1vPqplDv5rrFu5LrxTHbET0FiePTGi8zE"),
(
"./data/sparse_reconstruction_and_nerf_data/Friends-monica_apartment-disk.zip",
"1yNpU4M44gIuWEaC5tSAOKy-A6KX8wTm8",
),
(
"./data/sparse_reconstruction_and_nerf_data/HIMYM-red_apartment-disk.zip",
"1m76i46YYlpDKvlx-Jn7qDDAvkOVVVmqF",
),
(
"./data/sparse_reconstruction_and_nerf_data/Seinfeld-jerry_living_room-disk.zip",
"1YFoV3Gd9asKsRarZvzRTYtri1755rBPx",
),
("./data/sparse_reconstruction_and_nerf_data/TAAHM-kitchen-disk.zip", "1QpF8kVGfgkqm_cDY5sY_E82BT8VOWL38"),
(
"./data/sparse_reconstruction_and_nerf_data/TBBT-big_living_room-disk.zip",
"175BjkpxMAOt75ZVIjYpArxWpO8uRexgJ",
),
]
# !!! Large files !!! Only uncomment if you need the disk correspondences.
# data_filename_to_file_id += disk_data_filename_to_file_id
for data_filename, file_id in data_filename_to_file_id:
# Download the files
url = f"https://drive.google.com/uc?id={file_id}"
gdown.download(url, output=data_filename, quiet=False)
with zipfile.ZipFile(data_filename, "r") as zip_ref:
zip_ref.extractall(data_filename.replace(".zip", "").replace("-disk", ""))
Path(data_filename).unlink(missing_ok=True)
def download_blender():
"""Download blender dataset and format accordingly"""
url = "https://drive.google.com/uc?id=18JxhpWD-4ZmuFKLzKlAw-w5PpzZxXOcG"
download_path = Path("data/blender_data.zip")
gdown.download(url, output=str(download_path), quiet=False)
with zipfile.ZipFile(download_path, "r") as zip_ref:
zip_ref.extractall("data/")
unzip_path = Path("data/nerf_synthetic")
final_path = Path("data/blender")
unzip_path.rename(final_path)
download_path.unlink(missing_ok=True)
def main(dataset: Literal["blender", "friends"]):
"""Main download script to download all data"""
save_dir = Path("data/")
save_dir.mkdir(parents=True, exist_ok=True)
if dataset == "blender":
download_blender()
if dataset == "friends":
download_friends()
if __name__ == "__main__":
try:
# Requires dcargs>=0.3.3.
dcargs.extras.set_accent_color("bright_yellow")
except AttributeError:
pass
dcargs.cli(main)