-
Notifications
You must be signed in to change notification settings - Fork 7
/
utils.py
73 lines (56 loc) · 2.1 KB
/
utils.py
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from __future__ import absolute_import, division, print_function
import os
import hashlib
import zipfile
from six.moves import urllib
def readlines(filename):
"""Read all the lines in a text file and return as a list
"""
with open(filename, 'r') as f:
lines = f.read().splitlines()
return lines
def normalize_image(x):
"""Rescale image pixels to span range [0, 1]
"""
ma = float(x.max().cpu().data)
mi = float(x.min().cpu().data)
d = ma - mi if ma != mi else 1e5
return (x - mi) / d
def sec_to_hm(t):
"""Convert time in seconds to time in hours, minutes and seconds
e.g. 10239 -> (2, 50, 39)
"""
t = int(t)
s = t % 60
t //= 60
m = t % 60
t //= 60
return t, m, s
def sec_to_hm_str(t):
"""Convert time in seconds to a nice string
e.g. 10239 -> '02h50m39s'
"""
h, m, s = sec_to_hm(t)
return "{:02d}h{:02d}m{:02d}s".format(h, m, s)
if not os.path.exists("models"):
os.makedirs("models")
model_path = os.path.join("models", model_name)
def check_file_matches_md5(checksum, fpath):
if not os.path.exists(fpath):
return False
with open(fpath, 'rb') as f:
current_md5checksum = hashlib.md5(f.read()).hexdigest()
return current_md5checksum == checksum
# see if we have the model already downloaded...
if not os.path.exists(os.path.join(model_path, "encoder.pth")):
model_url, required_md5checksum = download_paths[model_name]
if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
print("-> Downloading pretrained model to {}".format(model_path + ".zip"))
urllib.request.urlretrieve(model_url, model_path + ".zip")
if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
print(" Failed to download a file which matches the checksum - quitting")
quit()
print(" Unzipping model...")
with zipfile.ZipFile(model_path + ".zip", 'r') as f:
f.extractall(model_path)
print(" Model unzipped to {}".format(model_path))