| import os |
| import os.path |
| import hashlib |
| import errno |
| from tqdm import tqdm |
|
|
|
|
| def gen_bar_updater(pbar): |
| def bar_update(count, block_size, total_size): |
| if pbar.total is None and total_size: |
| pbar.total = total_size |
| progress_bytes = count * block_size |
| pbar.update(progress_bytes - pbar.n) |
|
|
| return bar_update |
|
|
|
|
| def check_integrity(fpath, md5=None): |
| if md5 is None: |
| return True |
| if not os.path.isfile(fpath): |
| return False |
| md5o = hashlib.md5() |
| with open(fpath, 'rb') as f: |
| |
| for chunk in iter(lambda: f.read(1024 * 1024), b''): |
| md5o.update(chunk) |
| md5c = md5o.hexdigest() |
| if md5c != md5: |
| return False |
| return True |
|
|
|
|
| def makedir_exist_ok(dirpath): |
| """ |
| Python2 support for os.makedirs(.., exist_ok=True) |
| """ |
| try: |
| os.makedirs(dirpath) |
| except OSError as e: |
| if e.errno == errno.EEXIST: |
| pass |
| else: |
| raise |
|
|
|
|
| def download_url(url, root, filename=None, md5=None): |
| """Download a file from a url and place it in root. |
| Args: |
| url (str): URL to download file from |
| root (str): Directory to place downloaded file in |
| filename (str): Name to save the file under. If None, use the basename of the URL |
| md5 (str): MD5 checksum of the download. If None, do not check |
| """ |
| from six.moves import urllib |
|
|
| root = os.path.expanduser(root) |
| if not filename: |
| filename = os.path.basename(url) |
| fpath = os.path.join(root, filename) |
|
|
| makedir_exist_ok(root) |
|
|
| |
| if os.path.isfile(fpath) and check_integrity(fpath, md5): |
| print('Using downloaded and verified file: ' + fpath) |
| else: |
| try: |
| print('Downloading ' + url + ' to ' + fpath) |
| urllib.request.urlretrieve( |
| url, fpath, |
| reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True)) |
| ) |
| except OSError: |
| if url[:5] == 'https': |
| url = url.replace('https:', 'http:') |
| print('Failed download. Trying https -> http instead.' |
| ' Downloading ' + url + ' to ' + fpath) |
| urllib.request.urlretrieve( |
| url, fpath, |
| reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True)) |
| ) |
|
|
|
|
| def list_dir(root, prefix=False): |
| """List all directories at a given root |
| Args: |
| root (str): Path to directory whose folders need to be listed |
| prefix (bool, optional): If true, prepends the path to each result, otherwise |
| only returns the name of the directories found |
| """ |
| root = os.path.expanduser(root) |
| directories = list( |
| filter( |
| lambda p: os.path.isdir(os.path.join(root, p)), |
| os.listdir(root) |
| ) |
| ) |
|
|
| if prefix is True: |
| directories = [os.path.join(root, d) for d in directories] |
|
|
| return directories |
|
|
|
|
| def list_files(root, suffix, prefix=False): |
| """List all files ending with a suffix at a given root |
| Args: |
| root (str): Path to directory whose folders need to be listed |
| suffix (str or tuple): Suffix of the files to match, e.g. '.png' or ('.jpg', '.png'). |
| It uses the Python "str.endswith" method and is passed directly |
| prefix (bool, optional): If true, prepends the path to each result, otherwise |
| only returns the name of the files found |
| """ |
| root = os.path.expanduser(root) |
| files = list( |
| filter( |
| lambda p: os.path.isfile(os.path.join(root, p)) and p.endswith(suffix), |
| os.listdir(root) |
| ) |
| ) |
|
|
| if prefix is True: |
| files = [os.path.join(root, d) for d in files] |
|
|
| return files |
