Skip to content

Commit

Permalink
Make downloading more windows-friendly
Browse files Browse the repository at this point in the history
  • Loading branch information
VasuAgrawal committed Jan 23, 2024
1 parent 8952fb5 commit 5faada3
Showing 1 changed file with 26 additions and 23 deletions.
49 changes: 26 additions & 23 deletions nerfstudio/scripts/downloads/download_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import json
import os
import shutil
import subprocess
import tarfile
import zipfile
from dataclasses import dataclass
Expand Down Expand Up @@ -143,10 +144,10 @@ def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str,
file_id_or_zip_url = capture_name_to_file_id[capture_name]
if file_id_or_zip_url.endswith(".zip"):
url = file_id_or_zip_url # zip url
target_path = str(save_dir / f"{dataset_name}/{capture_name}")
target_path = str(save_dir / f"{dataset_name}" / f"{capture_name}")
os.makedirs(target_path, exist_ok=True)
download_path = Path(f"{target_path}.zip")
tmp_path = str(save_dir / ".temp")
tmp_path = str(save_dir / f"{dataset_name}" / f".temp_{capture_name}")
shutil.rmtree(tmp_path, ignore_errors=True)
os.makedirs(tmp_path, exist_ok=True)
try:
Expand All @@ -156,10 +157,10 @@ def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str,
run_command(f"wget {url} -O {download_path}", verbose=True)
else:
url = f"https://drive.google.com/uc?id={file_id_or_zip_url}" # file id
target_path = str(save_dir / f"{dataset_name}/{capture_name}")
target_path = str(save_dir / f"{dataset_name}" / f"{capture_name}")
os.makedirs(target_path, exist_ok=True)
download_path = Path(f"{target_path}.zip")
tmp_path = str(save_dir / ".temp")
tmp_path = str(save_dir / f"{dataset_name}" / f".temp_{capture_name}")
shutil.rmtree(tmp_path, ignore_errors=True)
os.makedirs(tmp_path, exist_ok=True)
try:
Expand All @@ -170,7 +171,7 @@ def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str,
with zipfile.ZipFile(download_path, "r") as zip_ref:
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."
assert len(inner_folders) == 1, f"There is more than one folder inside this zip file: {inner_folders}"
folder = os.path.join(tmp_path, inner_folders[0])
shutil.rmtree(target_path)
shutil.move(folder, target_path)
Expand Down Expand Up @@ -239,7 +240,9 @@ def download(self, save_dir: Path):
if os.path.exists(final_path):
shutil.rmtree(str(final_path))
download_path = save_dir / "dnerf_data.zip"
os.system(f"curl -L https://www.dropbox.com/s/raw/0bf6fl0ye2vz3vr/data.zip > {download_path}")
subprocess.run(
["curl", "-L", "https://www.dropbox.com/s/raw/0bf6fl0ye2vz3vr/data.zip", "-o", download_path], check=True
)
with zipfile.ZipFile(download_path, "r") as zip_ref:
zip_ref.extractall(str(save_dir))
unzip_path = save_dir / Path("data")
Expand Down Expand Up @@ -289,20 +292,20 @@ def download(self, save_dir: Path):
self.capture_name in phototourism_downloads
), f"Capture name {self.capture_name} not found in {phototourism_downloads.keys()}"
url = phototourism_downloads[self.capture_name]
target_path = str(save_dir / f"phototourism/{self.capture_name}")
target_path = str(save_dir / "phototourism" / f"{self.capture_name}")
os.makedirs(target_path, exist_ok=True)
download_path = Path(f"{target_path}.tar.gz")
tmp_path = str(save_dir / ".temp")
tmp_path = str(save_dir / "phototourism" / f".temp_{self.capture_name}")
shutil.rmtree(tmp_path, ignore_errors=True)
os.makedirs(tmp_path, exist_ok=True)

os.system(f"curl -L {url} > {download_path}")
subprocess.run(["curl", "-L", url, "-o", download_path], check=True)

with tarfile.open(download_path, "r:gz") as tar_ref:
tar_ref.extractall(str(tmp_path))

inner_folders = os.listdir(tmp_path)
assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
assert len(inner_folders) == 1, f"There is more than one folder inside this zip file: {inner_folders}"
folder = os.path.join(tmp_path, inner_folders[0])
shutil.rmtree(target_path)
shutil.move(folder, target_path)
Expand Down Expand Up @@ -337,7 +340,7 @@ class SDFstudioDemoDownload(DatasetDownload):
dataset_name: SDFstudioCaptureName = "sdfstudio-demo-data"

def download(self, save_dir: Path):
"""Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF)."""
"""Download the sdfstudio dataset (https://autonomousvision.github.io/sdfstudio/)."""
# TODO: give this code the same structure as download_nerfstudio

if self.dataset_name == "all":
Expand All @@ -352,17 +355,17 @@ def download(self, save_dir: Path):

url = sdfstudio_downloads[self.dataset_name]

target_path = str(save_dir / self.dataset_name)
target_path = str(save_dir / "sdfstudio" / self.dataset_name)
os.makedirs(target_path, exist_ok=True)

file_format = url[-4:]

download_path = Path(f"{target_path}{file_format}")
tmp_path = str(save_dir / ".temp")
tmp_path = str(save_dir / "sdfstudio" / f".temp_{self.dataset_name}")
shutil.rmtree(tmp_path, ignore_errors=True)
os.makedirs(tmp_path, exist_ok=True)

os.system(f"curl -L {url} > {download_path}")
subprocess.run(["curl", "-L", url, "-o", download_path], check=True)
if file_format == ".tar":
with tarfile.open(download_path, "r") as tar_ref:
tar_ref.extractall(str(tmp_path))
Expand All @@ -374,7 +377,7 @@ def download(self, save_dir: Path):
raise NotImplementedError

inner_folders = os.listdir(tmp_path)
assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
assert len(inner_folders) == 1, f"There is more than one folder inside this zip file: {inner_folders}"
folder = os.path.join(tmp_path, inner_folders[0])
shutil.rmtree(target_path)
shutil.move(folder, target_path)
Expand Down Expand Up @@ -421,21 +424,21 @@ def download(self, save_dir: Path):
self.capture_name in nerfosr_downloads
), f"Capture name {self.capture_name} not found in {nerfosr_downloads.keys()}"
url = nerfosr_downloads[self.capture_name]
target_path = str(save_dir / f"NeRF-OSR/Data/{self.capture_name}")
target_path = str(save_dir / "nerfosr" / f"{self.capture_name}")
os.makedirs(target_path, exist_ok=True)
download_path = Path(f"{target_path}.zip")
tmp_path = str(save_dir / ".temp")
tmp_path = str(save_dir / "nerfosr" / f".temp_{self.capture_name}")
shutil.rmtree(tmp_path, ignore_errors=True)
os.makedirs(tmp_path, exist_ok=True)

os.system(f"curl -L '{url}' > {download_path}")
subprocess.run(["curl", "-L", url, "-o", download_path], check=True)

# Extract the zip file
with zipfile.ZipFile(download_path, "r") as zip_ref:
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."
assert len(inner_folders) == 1, f"There is more than one folder inside this zip file: {inner_folders}"
folder = os.path.join(tmp_path, inner_folders[0])
shutil.rmtree(target_path)
shutil.move(folder, target_path)
Expand Down Expand Up @@ -475,20 +478,20 @@ def download(self, save_dir: Path) -> None:
self.capture_name in mill19_downloads
), f"Capture name {self.capture_name} not found in {mill19_downloads.keys()}"
url = mill19_downloads[self.capture_name]
target_path = save_dir / f"mill19/{self.capture_name}"
target_path = save_dir / "mill19" / f"{self.capture_name}"
target_path.mkdir(parents=True, exist_ok=True)
download_path = Path(f"{target_path}.tgz")
tmp_path = save_dir / ".temp"
tmp_path = save_dir / "mill19" / f".temp_{self.capture_name}"
shutil.rmtree(tmp_path, ignore_errors=True)
tmp_path.mkdir(parents=True, exist_ok=True)

os.system(f"curl -L {url} > {download_path}")
subprocess.run(["curl", "-L", url, "-o", download_path], check=True)

with tarfile.open(download_path, "r:gz") as tar_ref:
tar_ref.extractall(tmp_path)

inner_folders = list(tmp_path.iterdir())
assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
assert len(inner_folders) == 1, f"There is more than one folder inside this zip file: {inner_folders}"
folder = inner_folders[0]
shutil.rmtree(target_path)
folder.rename(target_path)
Expand Down

0 comments on commit 5faada3

Please sign in to comment.