utils.py 3.13 KB
Newer Older
soumith's avatar
soumith committed
1
import os
soumith's avatar
soumith committed
2
import os.path
soumith's avatar
soumith committed
3
4
import hashlib
import errno
5
6
7
from tqdm import tqdm


Holger Kohr's avatar
Holger Kohr committed
8
def gen_bar_updater(pbar):
9
    def bar_update(count, block_size, total_size):
Holger Kohr's avatar
Holger Kohr committed
10
11
12
13
        if pbar.total is None and total_size:
            pbar.total = total_size
        progress_bytes = count * block_size
        pbar.update(progress_bytes - pbar.n)
14
15

    return bar_update
soumith's avatar
soumith committed
16

soumith's avatar
soumith committed
17

18
19
20
def check_integrity(fpath, md5):
    if not os.path.isfile(fpath):
        return False
soumith's avatar
soumith committed
21
    md5o = hashlib.md5()
soumith's avatar
soumith committed
22
    with open(fpath, 'rb') as f:
soumith's avatar
soumith committed
23
        # read in 1MB chunks
24
        for chunk in iter(lambda: f.read(1024 * 1024), b''):
soumith's avatar
soumith committed
25
26
            md5o.update(chunk)
    md5c = md5o.hexdigest()
27
28
29
30
31
    if md5c != md5:
        return False
    return True


soumith's avatar
soumith committed
32
def download_url(url, root, filename, md5):
33
34
    from six.moves import urllib

35
    root = os.path.expanduser(root)
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    fpath = os.path.join(root, filename)

    try:
        os.makedirs(root)
    except OSError as e:
        if e.errno == errno.EEXIST:
            pass
        else:
            raise

    # downloads file
    if os.path.isfile(fpath) and check_integrity(fpath, md5):
        print('Using downloaded and verified file: ' + fpath)
    else:
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
50
51
        try:
            print('Downloading ' + url + ' to ' + fpath)
Holger Kohr's avatar
Holger Kohr committed
52
53
54
55
            urllib.request.urlretrieve(
                url, fpath,
                reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True))
            )
Tzu-Wei Huang's avatar
Tzu-Wei Huang committed
56
57
58
59
60
        except:
            if url[:5] == 'https':
                url = url.replace('https:', 'http:')
                print('Failed download. Trying https -> http instead.'
                      ' Downloading ' + url + ' to ' + fpath)
Holger Kohr's avatar
Holger Kohr committed
61
62
63
64
                urllib.request.urlretrieve(
                    url, fpath,
                    reporthook=gen_bar_updater(tqdm(unit='B', unit_scale=True))
                )
Sanyam Kapoor's avatar
Sanyam Kapoor committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110


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