-
Notifications
You must be signed in to change notification settings - Fork 3
/
utils.py
165 lines (149 loc) · 5.03 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
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import glob
import os
from PIL import Image
import urllib.request
from tqdm import tqdm
import numpy as np
import PIL.Image
import sys
from io import BytesIO
import IPython.display
import numpy as np
from math import ceil
from PIL import Image, ImageDraw
import os
from IPython.display import HTML
from base64 import b64encode
import imageio
def show_animation(movie_name):
mp4 = open(movie_name,'rb').read()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
return HTML("""
<video width=400 controls>
<source src="%s" type="video/mp4">
</video>
""" % data_url)
def imshow(a, format='png', jpeg_fallback=True):
a = np.asarray(a, dtype=np.uint8)
str_file = BytesIO()
PIL.Image.fromarray(a).save(str_file, format)
im_data = str_file.getvalue()
try:
disp = IPython.display.display(IPython.display.Image(im_data))
except IOError:
if jpeg_fallback and format != 'jpeg':
print ('Warning: image was too large to display in format "{}"; '
'trying jpeg instead.').format(format)
return imshow(a, format='jpeg')
else:
raise
return disp
def clamp(x, minimum, maximum):
return max(minimum, min(x, maximum))
def create_image_grid(images, scale=0.25, rows=1):
w,h = images[0].size
w = int(w*scale)
h = int(h*scale)
height = rows*h
cols = ceil(len(images) / rows)
width = cols*w
canvas = PIL.Image.new('RGBA', (width,height), 'white')
for i,img in enumerate(images):
img = img.resize((w,h), PIL.Image.ANTIALIAS)
canvas.paste(img, (w*(i % cols), h*(i // cols)))
return canvas
def find_latest_pkl(path):
curr_best = 0
latest_pkl = ''
for pkl in glob.glob(f'{path}/*.pkl'):
ckpt_number = int(pkl.split('-')[-1][:-4])
if curr_best < ckpt_number:
curr_best = ckpt_number
latest_pkl = pkl
return latest_pkl
def resize(path, dim = (512, 512)):
dirs = os.listdir(path)
out_path = f'{path}/{dim[0]}x{dim[1]}'
os.makedirs(out_path, exist_ok=True)
for item in log_progress(dirs):
img_path = f'{path}/{item}'
if os.path.isfile(img_path):
im = Image.open(img_path)
imResize = im.resize(dim, Image.ANTIALIAS).convert('RGB')
imResize.save(f'{out_path}/{item}', 'JPEG', quality=90)
return out_path
def resize_dirs(path, out_dir, dim = (512, 512)):
sub_dirs = os.listdir(path)
for sub_dir in sub_dirs:
out_path = f'{out_dir}/{sub_dir}'
os.makedirs(out_path, exist_ok=True)
for item in log_progress(os.listdir(f'{path}/{sub_dir}/')[:10]):
img_path = f'{path}/{sub_dir}/{item}'
if os.path.isfile(img_path):
im = Image.open(img_path)
imResize = im.resize(dim, Image.ANTIALIAS).convert('RGB')
imResize.save(f'{out_path}/{item}', 'JPEG', quality=90)
return out_dir
class DownloadProgressBar(tqdm):
def update_to(self, b=1, bsize=1, tsize=None):
if tsize is not None:
self.total = tsize
self.update(b * bsize - self.n)
# https://stackoverflow.com/a/53877507
def download_url(url, output_path):
with DownloadProgressBar(unit='B', unit_scale=True,
miniters=1, desc=url.split('/')[-1]) as t:
urllib.request.urlretrieve(url, filename=output_path, reporthook=t.update_to)
# Taken from https://github.com/alexanderkuk/log-progress
def log_progress(sequence, every=1, size=None, name='Items'):
from ipywidgets import IntProgress, HTML, VBox
from IPython.display import display
is_iterator = False
if size is None:
try:
size = len(sequence)
except TypeError:
is_iterator = True
if size is not None:
if every is None:
if size <= 200:
every = 1
else:
every = int(size / 200) # every 0.5%
else:
assert every is not None, 'sequence is iterator, set every'
if is_iterator:
progress = IntProgress(min=0, max=1, value=1)
progress.bar_style = 'info'
else:
progress = IntProgress(min=0, max=size, value=0)
label = HTML()
box = VBox(children=[label, progress])
display(box)
index = 0
try:
for index, record in enumerate(sequence, 1):
if index == 1 or index % every == 0:
if is_iterator:
label.value = '{name}: {index} / ?'.format(
name=name,
index=index
)
else:
progress.value = index
label.value = u'{name}: {index} / {size}'.format(
name=name,
index=index,
size=size
)
yield record
except:
progress.bar_style = 'danger'
raise
else:
progress.bar_style = 'success'
progress.value = index
label.value = "{name}: {index}".format(
name=name,
index=str(index or '?')
)