Spaces:
Running
Running
File size: 1,707 Bytes
79230dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import os
import requests
import zipfile
from tqdm import tqdm
def download_file(url: str, save_dir='./', overwrite=False, unzip=True):
os.makedirs(save_dir, exist_ok=True)
file_name = url.split('/')[-1]
file_path = os.path.join(save_dir, file_name)
if os.path.exists(file_path) and not overwrite:
pass
else:
print('Downloading file {} from {}...'.format(file_path, url))
r = requests.get(url, stream=True)
print(r.status_code)
if r.status_code != 200:
raise RuntimeError('Failed downloading url {}!'.format(url))
total_length = r.headers.get('content-length')
with open(file_path, 'wb') as f:
if total_length is None: # no content length header
for chunk in r.iter_content(chunk_size=1024):
if chunk: # filter out keep-alive new chunks
f.write(chunk)
else:
total_length = int(total_length)
print('file length: ', int(total_length / 1024. + 0.5))
for chunk in tqdm(r.iter_content(chunk_size=1024),
total=int(total_length / 1024. + 0.5),
unit='KB',
unit_scale=False,
dynamic_ncols=True):
f.write(chunk)
if unzip and file_path.endswith('.zip'):
save_dir = file_path.split('.')[0]
if os.path.isdir(save_dir) and os.path.exists(save_dir):
pass
else:
with zipfile.ZipFile(file_path, 'r') as zip_ref:
zip_ref.extractall(save_dir)
return save_dir, file_path
|