def get_loss_history_preview(loss_history, iter, w, c):
loss_history = np.array(loss_history.copy())
lh_height = 100
lh_img = np.ones((lh_height, w, c)) * 0.1
if len(loss_history) != 0:
loss_count = len(loss_history[0])
lh_len = len(loss_history)
l_per_col = lh_len / w
plist_max = [[
max(
0.0, loss_history[int(col * l_per_col)][p], *[
loss_history[i_ab][p] for i_ab in range(
int(col * l_per_col), int((col + 1) * l_per_col))
]) for p in range(loss_count)
] for col in range(w)]
plist_min = [[
min(
plist_max[col][p], loss_history[int(col * l_per_col)][p],
*[
loss_history[i_ab][p] for i_ab in range(
int(col * l_per_col), int((col + 1) * l_per_col))
]) for p in range(loss_count)
] for col in range(w)]
plist_abs_max = np.mean(loss_history[len(loss_history) // 5:]) * 2
for col in range(0, w):
for p in range(0, loss_count):
point_color = [1.0] * c
point_color[0:3] = colorsys.hsv_to_rgb(
p * (1.0 / loss_count), 1.0, 1.0)
ph_max = int(
(plist_max[col][p] / plist_abs_max) * (lh_height - 1))
ph_max = np.clip(ph_max, 0, lh_height - 1)
ph_min = int(
(plist_min[col][p] / plist_abs_max) * (lh_height - 1))
ph_min = np.clip(ph_min, 0, lh_height - 1)
for ph in range(ph_min, ph_max + 1):
lh_img[(lh_height - ph - 1), col] = point_color
lh_lines = 5
lh_line_height = (lh_height - 1) / lh_lines
for i in range(0, lh_lines + 1):
lh_img[int(i * lh_line_height), :] = (0.8, ) * c
last_line_t = int((lh_lines - 1) * lh_line_height)
last_line_b = int(lh_lines * lh_line_height)
lh_text = 'Iter: %d' % (iter) if iter != 0 else ''
lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image(
(last_line_b - last_line_t, w, c), lh_text, color=[0.8] * c)
return lh_img
def get_screen_status_block(self, w, c):
if self.screen_status_block_dirty:
self.screen_status_block_dirty = False
lines = [
'Polys current/max = %d/%d' % (self.ie_polys.n, self.ie_polys.n_max),
]
if self.get_status_lines_func is not None:
lines += self.get_status_lines_func()
lines_count = len(lines)
h_line = 21
h = lines_count * h_line
img = np.ones ( (h,w,c) ) * 0.1
for i in range(lines_count):
img[ i*h_line:(i+1)*h_line, 0:w] += \
imagelib.get_text_image ( (h_line,w,c), lines[i], color=[0.8]*c )
self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8)
return self.screen_status_block
def create_preview_pane_image(previews, selected_preview, loss_history,
show_last_history_iters_count, iteration, batch_size, zoom=Zoom.ZOOM_100):
scaled_previews = scale_previews(previews, zoom)
selected_preview_name = scaled_previews[selected_preview][0]
selected_preview_rgb = scaled_previews[selected_preview][1]
h, w, c = selected_preview_rgb.shape
# HEAD
head_lines = [
'[s]:save [enter]:exit [-/+]:zoom: %s' % zoom.label,
'[p]:update [space]:next preview [l]:change history range',
'Preview: "%s" [%d/%d]' % (selected_preview_name, selected_preview + 1, len(previews))
]
head_line_height = int(15 * zoom.scale)
head_height = len(head_lines) * head_line_height
head = np.ones((head_height, w, c)) * 0.1
for i in range(0, len(head_lines)):
t = i * head_line_height
b = (i + 1) * head_line_height
head[t:b, 0:w] += imagelib.get_text_image((head_line_height, w, c), head_lines[i], color=[0.8] * c)
final = head
if loss_history is not None:
if show_last_history_iters_count == 0:
loss_history_to_show = loss_history
else:
loss_history_to_show = loss_history[-show_last_history_iters_count:]
lh_height = int(100 * zoom.scale)
lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iteration, w, c, lh_height)
final = np.concatenate([final, lh_img], axis=0)
final = np.concatenate([final, selected_preview_rgb], axis=0)
final = np.clip(final, 0, 1)
return (final * 255).astype(np.uint8)
def main(**kwargs):
io.log_info("Running trainer.\r\n")
no_preview = kwargs.get('no_preview', False)
s2c = queue.Queue()
c2s = queue.Queue()
e = threading.Event()
thread = threading.Thread(target=trainerThread,
args=(s2c, c2s, e),
kwargs=kwargs)
thread.start()
e.wait() #Wait for inital load to occur.
if no_preview:
while True:
if not c2s.empty():
input = c2s.get()
op = input.get('op', '')
if op == 'close':
break
try:
io.process_messages(0.1)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
else:
wnd_name = "Training preview"
io.named_window(wnd_name)
io.capture_keys(wnd_name)
previews = None
loss_history = None
selected_preview = 0
update_preview = False
is_showing = False
is_waiting_preview = False
show_last_history_iters_count = 0
iter = 0
while True:
if not c2s.empty():
input = c2s.get()
op = input['op']
if op == 'show':
is_waiting_preview = False
loss_history = input[
'loss_history'] if 'loss_history' in input.keys(
) else None
previews = input['previews'] if 'previews' in input.keys(
) else None
iter = input['iter'] if 'iter' in input.keys() else 0
if previews is not None:
max_w = 0
max_h = 0
for (preview_name, preview_rgb) in previews:
(h, w, c) = preview_rgb.shape
max_h = max(max_h, h)
max_w = max(max_w, w)
max_size = 800
if max_h > max_size:
max_w = int(max_w / (max_h / max_size))
max_h = max_size
#make all previews size equal
for preview in previews[:]:
(preview_name, preview_rgb) = preview
(h, w, c) = preview_rgb.shape
if h != max_h or w != max_w:
previews.remove(preview)
previews.append(
(preview_name,
cv2.resize(preview_rgb, (max_w, max_h))))
selected_preview = selected_preview % len(previews)
update_preview = True
elif op == 'close':
break
if update_preview:
update_preview = False
selected_preview_name = previews[selected_preview][0]
selected_preview_rgb = previews[selected_preview][1]
(h, w, c) = selected_preview_rgb.shape
# HEAD
head_lines = [
'[s]:save [enter]:exit',
'[p]:update [space]:next preview [l]:change history range',
'Preview: "%s" [%d/%d]' %
(selected_preview_name, selected_preview + 1,
len(previews))
]
head_line_height = 15
head_height = len(head_lines) * head_line_height
head = np.ones((head_height, w, c)) * 0.1
for i in range(0, len(head_lines)):
t = i * head_line_height
b = (i + 1) * head_line_height
head[t:b, 0:w] += imagelib.get_text_image(
(head_line_height, w, c),
head_lines[i],
color=[0.8] * c)
final = head
if loss_history is not None:
if show_last_history_iters_count == 0:
loss_history_to_show = loss_history
else:
loss_history_to_show = loss_history[
-show_last_history_iters_count:]
lh_img = models.ModelBase.get_loss_history_preview(
loss_history_to_show, iter, w, c)
final = np.concatenate([final, lh_img], axis=0)
final = np.concatenate([final, selected_preview_rgb], axis=0)
final = np.clip(final, 0, 1)
io.show_image(wnd_name, (final * 255).astype(np.uint8))
is_showing = True
key_events = io.get_key_events(wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[
-1] if len(key_events) > 0 else (0, 0, False, False, False)
if key == ord('\n') or key == ord('\r'):
s2c.put({'op': 'close'})
elif key == ord('s'):
s2c.put({'op': 'save'})
elif key == ord('p'):
if not is_waiting_preview:
is_waiting_preview = True
s2c.put({'op': 'preview'})
elif key == ord('l'):
if show_last_history_iters_count == 0:
show_last_history_iters_count = 5000
elif show_last_history_iters_count == 5000:
show_last_history_iters_count = 10000
elif show_last_history_iters_count == 10000:
show_last_history_iters_count = 50000
elif show_last_history_iters_count == 50000:
show_last_history_iters_count = 100000
elif show_last_history_iters_count == 100000:
show_last_history_iters_count = 0
update_preview = True
elif key == ord(' '):
selected_preview = (selected_preview + 1) % len(previews)
update_preview = True
try:
io.process_messages(0.1)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
io.destroy_all_windows()
def main(**kwargs):
io.log_info("Running trainer.\r\n")
no_preview = kwargs.get('no_preview', False)
flask_preview = kwargs.get('flask_preview', False)
s2c = queue.Queue()
c2s = queue.Queue()
e = threading.Event()
previews = None
loss_history = None
selected_preview = 0
update_preview = False
is_waiting_preview = False
show_last_history_iters_count = 0
iteration = 0
batch_size = 1
zoom = Zoom.ZOOM_100
if flask_preview:
from flaskr.app import create_flask_app
s2flask = queue.Queue()
socketio, flask_app = create_flask_app(s2c, c2s, s2flask, kwargs)
thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e, socketio), kwargs=kwargs)
thread.start()
e.wait() # Wait for inital load to occur.
flask_t = threading.Thread(target=socketio.run, args=(flask_app,),
kwargs={'debug': True, 'use_reloader': False})
flask_t.start()
while True:
if not c2s.empty():
item = c2s.get()
op = item['op']
if op == 'show':
is_waiting_preview = False
loss_history = item['loss_history'] if 'loss_history' in item.keys() else None
previews = item['previews'] if 'previews' in item.keys() else None
iteration = item['iter'] if 'iter' in item.keys() else 0
# batch_size = input['batch_size'] if 'iter' in input.keys() else 1
if previews is not None:
update_preview = True
elif op == 'update':
if not is_waiting_preview:
is_waiting_preview = True
s2c.put({'op': 'preview'})
elif op == 'next_preview':
selected_preview = (selected_preview + 1) % len(previews)
update_preview = True
elif op == 'change_history_range':
if show_last_history_iters_count == 0:
show_last_history_iters_count = 5000
elif show_last_history_iters_count == 5000:
show_last_history_iters_count = 10000
elif show_last_history_iters_count == 10000:
show_last_history_iters_count = 50000
elif show_last_history_iters_count == 50000:
show_last_history_iters_count = 100000
elif show_last_history_iters_count == 100000:
show_last_history_iters_count = 0
update_preview = True
elif op == 'close':
s2c.put({'op': 'close'})
break
elif op == 'zoom_prev':
zoom = zoom.prev()
update_preview = True
elif op == 'zoom_next':
zoom = zoom.next()
update_preview = True
if update_preview:
update_preview = False
selected_preview = selected_preview % len(previews)
preview_pane_image = create_preview_pane_image(previews,
selected_preview,
loss_history,
show_last_history_iters_count,
iteration,
batch_size,
zoom)
# io.show_image(wnd_name, preview_pane_image)
model_path = Path(kwargs.get('saved_models_path', ''))
filename = 'preview.png'
preview_file = str(model_path / filename)
cv2.imwrite(preview_file, preview_pane_image)
s2flask.put({'op': 'show'})
socketio.emit('preview', {'iter': iteration, 'loss': loss_history[-1]})
try:
io.process_messages(0.01)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
else:
thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e), kwargs=kwargs)
thread.start()
e.wait() # Wait for inital load to occur.
if no_preview:
while True:
if not c2s.empty():
item = c2s.get()
op = item.get('op', '')
if op == 'close':
break
try:
io.process_messages(0.1)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
else:
wnd_name = "Training preview"
io.named_window(wnd_name)
io.capture_keys(wnd_name)
previews = None
loss_history = None
selected_preview = 0
update_preview = False
is_showing = False
is_waiting_preview = False
show_last_history_iters_count = 0
iter = 0
while True:
if not c2s.empty():
item = c2s.get()
op = item['op']
if op == 'show':
is_waiting_preview = False
loss_history = item['loss_history'] if 'loss_history' in item.keys() else None
previews = item['previews'] if 'previews' in item.keys() else None
iter = item['iter'] if 'iter' in item.keys() else 0
if previews is not None:
max_w = 0
max_h = 0
for (preview_name, preview_rgb) in previews:
(h, w, c) = preview_rgb.shape
max_h = max(max_h, h)
max_w = max(max_w, w)
max_size = 800
if max_h > max_size:
max_w = int(max_w / (max_h / max_size))
max_h = max_size
# make all previews size equal
for preview in previews[:]:
(preview_name, preview_rgb) = preview
(h, w, c) = preview_rgb.shape
if h != max_h or w != max_w:
previews.remove(preview)
previews.append((preview_name, cv2.resize(preview_rgb, (max_w, max_h))))
selected_preview = selected_preview % len(previews)
update_preview = True
elif op == 'close':
break
if update_preview:
update_preview = False
selected_preview_name = previews[selected_preview][0]
selected_preview_rgb = previews[selected_preview][1]
(h, w, c) = selected_preview_rgb.shape
# HEAD
head_lines = [
'[s]:save [b]:backup [enter]:exit',
'[p]:update [space]:next preview [l]:change history range',
'Preview: "%s" [%d/%d]' % (selected_preview_name, selected_preview + 1, len(previews))
]
head_line_height = 15
head_height = len(head_lines) * head_line_height
head = np.ones((head_height, w, c)) * 0.1
for i in range(0, len(head_lines)):
t = i * head_line_height
b = (i + 1) * head_line_height
head[t:b, 0:w] += imagelib.get_text_image((head_line_height, w, c), head_lines[i], color=[0.8] * c)
final = head
if loss_history is not None:
if show_last_history_iters_count == 0:
loss_history_to_show = loss_history
else:
loss_history_to_show = loss_history[-show_last_history_iters_count:]
lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iter, w, c)
final = np.concatenate([final, lh_img], axis=0)
final = np.concatenate([final, selected_preview_rgb], axis=0)
final = np.clip(final, 0, 1)
io.show_image(wnd_name, (final * 255).astype(np.uint8))
is_showing = True
key_events = io.get_key_events(wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (
0, 0, False, False, False)
if key == ord('\n') or key == ord('\r'):
s2c.put({'op': 'close'})
elif key == ord('s'):
s2c.put({'op': 'save'})
elif key == ord('b'):
s2c.put({'op': 'backup'})
elif key == ord('p'):
if not is_waiting_preview:
is_waiting_preview = True
s2c.put({'op': 'preview'})
elif key == ord('l'):
if show_last_history_iters_count == 0:
show_last_history_iters_count = 5000
elif show_last_history_iters_count == 5000:
show_last_history_iters_count = 10000
elif show_last_history_iters_count == 10000:
show_last_history_iters_count = 50000
elif show_last_history_iters_count == 50000:
show_last_history_iters_count = 100000
elif show_last_history_iters_count == 100000:
show_last_history_iters_count = 0
update_preview = True
elif key == ord(' '):
selected_preview = (selected_preview + 1) % len(previews)
update_preview = True
try:
io.process_messages(0.1)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
io.destroy_all_windows()