mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
* Set dcargs accent color * Bump required dcargs version * color -> colour * Bump dcargs to 0.3.6
85 lines
3.1 KiB
Python
Executable File
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)
|