John6666's picture
Upload 351 files
e84842d verified
raw
history blame
1.98 kB
"""
Copyright (c) 2022, salesforce.com, inc.
All rights reserved.
SPDX-License-Identifier: BSD-3-Clause
For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import os
from pathlib import Path
from omegaconf import OmegaConf
from lavis.common.utils import (
cleanup_dir,
download_and_extract_archive,
get_abs_path,
get_cache_path,
)
DATA_URL = "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/datasets/didemo/didemo_videos.tar.gz"
def download_datasets(root, url):
"""
Download the Imagenet-R dataset archives and expand them
in the folder provided as parameter
"""
download_and_extract_archive(url=url, download_root=root)
def move_files(download_path, storage_path):
"""
Move files from download_path to storage_path
"""
print("Moving to {}".format(storage_path))
os.makedirs(storage_path, exist_ok=True)
for file_name in os.listdir(download_path):
os.rename(
os.path.join(download_path, file_name),
os.path.join(storage_path, file_name),
)
if __name__ == "__main__":
config_path = get_abs_path("configs/datasets/didemo/defaults_ret.yaml")
storage_dir = OmegaConf.load(
config_path
).datasets.didemo_retrieval.build_info.videos.storage
download_dir = Path(get_cache_path(storage_dir)).parent / "download"
storage_dir = Path(get_cache_path(storage_dir))
if storage_dir.exists():
print(f"Dataset already exists at {storage_dir}. Aborting.")
exit(0)
try:
print("Downloading {} to {}".format(DATA_URL, download_dir))
download_datasets(download_dir, DATA_URL)
except Exception as e:
# remove download dir if failed
cleanup_dir(download_dir)
print("Failed to download or extracting datasets. Aborting.")
move_files(download_dir / "videos", storage_dir)
cleanup_dir(download_dir)