Add unzip flag to download() (#3002)

pull/3004/head
Glenn Jocher 2021-05-01 12:37:41 +02:00 committed by GitHub
parent f55730ee9c
commit 4d7bca7648
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -183,7 +183,7 @@ def check_dataset(dict):
raise Exception('Dataset not found.')
def download(url, dir='.', threads=1):
def download(url, dir='.', unzip=True, threads=1):
# Multi-threaded file download and unzip function
def download_one(url, dir):
# Download 1 file
@ -191,7 +191,7 @@ def download(url, dir='.', threads=1):
if not f.exists():
print(f'Downloading {url} to {f}...')
torch.hub.download_url_to_file(url, f, progress=True) # download
if f.suffix in ('.zip', '.gz'):
if unzip and f.suffix in ('.zip', '.gz'):
print(f'Unzipping {f}...')
if f.suffix == '.zip':
os.system(f'unzip -qo {f} -d {dir} && rm {f}') # unzip -quiet -overwrite