github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-22 22:34:25 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

better flask previews

Browse source
This commit is contained in:
jh 2019-09-13 16:23:07 -07:00
commit 637bf69eb7
1 changed files with 255 additions and 203 deletions
Download patch file Download diff file Expand all files Collapse all files

356
mainscripts/FlaskTrainer.py
View file

@ -3,8 +3,9 @@ import traceback
import queue import queue
import threading import threading
import time import time
from io import BytesIO from enum import Enum
import base64 from os.path import getmtime
import numpy as np import numpy as np
import itertools import itertools
from pathlib import Path from pathlib import Path
@ -14,7 +15,7 @@ import cv2
import models import models
from interact import interact as io from interact import interact as io
from flask import Flask, send_file, Response, render_template, render_template_string, request, g from flask import Flask, send_file, Response, render_template, render_template_string, request, g
from flask_caching import Cache # from flask_socketio import SocketIO
def trainerThread (s2c, c2s, e, args, device_args): def trainerThread (s2c, c2s, e, args, device_args):
while True: while True:
@ -57,7 +58,6 @@ def trainerThread(s2c, c2s, e, args, device_args):
shared_state = { 'after_save' : False } shared_state = { 'after_save' : False }
loss_string = "" loss_string = ""
save_iter = model.get_iter() save_iter = model.get_iter()
def model_save(): def model_save():
if not debug and not is_reached_goal: if not debug and not is_reached_goal:
io.log_info ("Saving....", end='\r') io.log_info ("Saving....", end='\r')
@ -67,13 +67,13 @@ def trainerThread(s2c, c2s, e, args, device_args):
def send_preview(): def send_preview():
if not debug: if not debug:
previews = model.get_previews() previews = model.get_previews()
c2s.put({'op': 'show', 'previews': previews, 'iter': model.get_iter(), c2s.put ( {'op':'show', 'previews': previews, 'iter':model.get_iter(), 'loss_history': model.get_loss_history().copy() } )
'loss_history': model.get_loss_history().copy()})
else: else:
previews = [( 'debug, press update for new', model.debug_one_iter())] previews = [( 'debug, press update for new', model.debug_one_iter())]
c2s.put ( {'op':'show', 'previews': previews} ) c2s.put ( {'op':'show', 'previews': previews} )
e.set() #Set the GUI Thread as Ready e.set() #Set the GUI Thread as Ready
if model.is_first_run(): if model.is_first_run():
model_save() model_save()
@ -81,8 +81,7 @@ def trainerThread(s2c, c2s, e, args, device_args):
if is_reached_goal: if is_reached_goal:
io.log_info('Model already trained to target iteration. You can use preview.') io.log_info('Model already trained to target iteration. You can use preview.')
else: else:
io.log_info('Starting. Target iteration: %d. Press "Enter" to stop training and save model.' % ( io.log_info('Starting. Target iteration: %d. Press "Enter" to stop training and save model.' % ( model.get_target_iter() ) )
model.get_target_iter()))
else: else:
io.log_info('Starting. Press "Enter" to stop training and save model.') io.log_info('Starting. Press "Enter" to stop training and save model.')
@ -116,12 +115,9 @@ def trainerThread(s2c, c2s, e, args, device_args):
loss_history = model.get_loss_history() loss_history = model.get_loss_history()
time_str = time.strftime("[%H:%M:%S]") time_str = time.strftime("[%H:%M:%S]")
if iter_time >= 10: if iter_time >= 10:
loss_string = "{0}[#{1:06d}][{2:.5s}s][bs: {3}]".format(time_str, iter, loss_string = "{0}[#{1:06d}][{2:.5s}s][bs: {3}]".format ( time_str, iter, '{:0.4f}'.format(iter_time), batch_size )
'{:0.4f}'.format(iter_time),
batch_size)
else: else:
loss_string = "{0}[#{1:06d}][{2:04d}ms][bs: {3}]".format(time_str, iter, loss_string = "{0}[#{1:06d}][{2:04d}ms][bs: {3}]".format ( time_str, iter, int(iter_time*1000), batch_size)
int(iter_time * 1000), batch_size)
if shared_state['after_save']: if shared_state['after_save']:
shared_state['after_save'] = False shared_state['after_save'] = False
@ -179,6 +175,8 @@ def trainerThread(s2c, c2s, e, args, device_args):
if i == -1: if i == -1:
break break
model.finalize() model.finalize()
except Exception as e: except Exception as e:
@ -188,71 +186,80 @@ def trainerThread(s2c, c2s, e, args, device_args):
c2s.put ( {'op':'close'} ) c2s.put ( {'op':'close'} )
class Preview: class Zoom(Enum):
def __init__(self, c2s, s2c, preview_queue): ZOOM_25 = (1/4, '25%')
self.c2s = c2s ZOOM_33 = (1/3, '33%')
self.s2c = s2c ZOOM_50 = (1/2, '50%')
self.preview_queue = preview_queue ZOOM_67 = (2/3, '67%')
# self.wnd_name = "Training preview" ZOOM_75 = (3/4, '75%')
# io.named_window(wnd_name) ZOOM_80 = (4/5, '80%')
# io.capture_keys(wnd_name) ZOOM_90 = (9/10, '90%')
ZOOM_100 = (1, '100%')
ZOOM_110 = (11/10, '110%')
ZOOM_125 = (5/4, '125%')
ZOOM_150 = (3/2, '150%')
ZOOM_175 = (7/4, '175%')
ZOOM_200 = (2, '200%')
ZOOM_250 = (5/2, '250%')
ZOOM_300 = (3, '300%')
ZOOM_400 = (4, '400%')
ZOOM_500 = (5, '500%')
self.previews = None def __init__(self, scale, label):
self.loss_history = None self.scale = scale
self.selected_preview = 0 self.label = label
self.update_preview = False
self.is_showing = False
self.is_waiting_preview = False
self.show_last_history_iters_count = 0
self.iter = 0
self.batch_size = 1
self.preview_min_height = 512
self.preview_max_height = 1024
self.close = False
def get_preview(self): def prev(self):
while not self.close: cls = self.__class__
self.process_queue_items() members = list(cls)
self.update_preview_frame() index = members.index(self) - 1
if index < 0:
return self
return members[index]
def process_queue_items(self): def next(self):
if not self.c2s.empty(): cls = self.__class__
input = self.c2s.get() members = list(cls)
op = input['op'] index = members.index(self) + 1
if op == 'show': if index >= len(members):
self.is_waiting_preview = False return self
self.loss_history = input['loss_history'] if 'loss_history' in input.keys() else None return members[index]
self.previews = input['previews'] if 'previews' in input.keys() else None
self.iter = input['iter'] if 'iter' in input.keys() else 0
if self.previews is not None:
self.resize_previews()
self.selected_preview = self.selected_preview % len(self.previews)
self.update_preview = True
elif op == 'close':
self.close = True
elif op == 'update':
self.update()
elif op == 'next_preview':
self.next_preview()
elif op == 'change_history_range':
self.change_history_range()
def update_preview_frame(self): def scale_previews(previews, zoom=Zoom.ZOOM_100):
if self.update_preview: scaled = []
self.update_preview = False for preview in previews:
preview_name, preview_rgb = preview
scale_factor = zoom.scale
if scale_factor < 1:
scaled.append((preview_name, cv2.resize(preview_rgb, (0, 0),
fx=scale_factor,
fy=scale_factor,
interpolation=cv2.INTER_AREA)))
elif scale_factor > 1:
scaled.append((preview_name, cv2.resize(preview_rgb, (0, 0),
fx=scale_factor,
fy=scale_factor,
interpolation=cv2.INTER_LANCZOS4)))
else:
scaled.append((preview_name, preview_rgb))
return scaled
selected_preview_name = self.previews[self.selected_preview][0]
selected_preview_rgb = self.previews[self.selected_preview][1] def create_preview_pane_image(previews, selected_preview, loss_history,
(h, w, c) = selected_preview_rgb.shape 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
head_lines = [ head_lines = [
'[s]:save [enter]:exit', '[s]:save [enter]:exit [-/+]:zoom: %s' % zoom.label,
'[p]:update [space]:next preview [l]:change history range', '[p]:update [space]:next preview [l]:change history range',
'Preview: "%s" [%d/%d]' % (selected_preview_name, self.selected_preview + 1, len(self.previews)) 'Preview: "%s" [%d/%d]' % (selected_preview_name,selected_preview+1, len(previews))
] ]
head_line_height = 15 head_line_height = int(15 * zoom.scale)
head_height = len(head_lines) * head_line_height head_height = len(head_lines) * head_line_height
head = np.ones((head_height, w, c)) * 0.1 head = np.ones((head_height, w, c)) * 0.1
@ -263,85 +270,32 @@ class Preview:
final = head final = head
if self.loss_history is not None: if loss_history is not None:
if self.show_last_history_iters_count == 0: if show_last_history_iters_count == 0:
loss_history_to_show = self.loss_history loss_history_to_show = loss_history
else: else:
loss_history_to_show = self.loss_history[-self.show_last_history_iters_count:] 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, self.iter, self.batch_size, w, lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iteration, batch_size, w, c, lh_height)
c)
final = np.concatenate ( [final, lh_img], axis=0 ) final = np.concatenate ( [final, lh_img], axis=0 )
final = np.concatenate([final, selected_preview_rgb], axis=0) final = np.concatenate([final, selected_preview_rgb], axis=0)
final = np.clip(final, 0, 1) final = np.clip(final, 0, 1)
preview_pane = (final * 255).astype(np.uint8) return (final*255).astype(np.uint8)
retval, buffer = cv2.imencode('.jpg', preview_pane)
# jpg_as_text = base64.b64encode(buffer)
jpg_as_text = buffer.tostring()
self.preview_queue.put(jpg_as_text)
def resize_previews(self):
preview_height = max((h for h, w, c in (im.shape for name, im in self.previews)))
if preview_height > self.preview_max_height:
preview_height = self.preview_max_height
elif preview_height < self.preview_min_height:
preview_height = self.preview_min_height
# make all previews size equal
for p in self.previews[:]:
(preview_name, preview_rgb) = p
(h, w, c) = preview_rgb.shape
if h != preview_height:
scale_factor = preview_height / float(h)
self.previews.remove(p)
self.previews.append((preview_name, cv2.resize(preview_rgb, (0, 0),
fx=scale_factor,
fy=scale_factor,
interpolation=cv2.INTER_AREA)))
self.selected_preview = self.selected_preview % len(self.previews)
def save(self):
self.s2c.put({'op': 'save'})
def exit(self):
self.s2c.put({'op': 'close'})
def update(self):
if not self.is_waiting_preview:
self.is_waiting_preview = True
self.s2c.put({'op': 'preview'})
def next_preview(self):
self.selected_preview = (self.selected_preview + 1) % len(self.previews)
self.update_preview = True
def change_history_range(self):
if self.show_last_history_iters_count == 0:
self.show_last_history_iters_count = 5000
elif self.show_last_history_iters_count == 5000:
self.show_last_history_iters_count = 10000
elif self.show_last_history_iters_count == 10000:
self.show_last_history_iters_count = 50000
elif self.show_last_history_iters_count == 50000:
self.show_last_history_iters_count = 100000
elif self.show_last_history_iters_count == 100000:
self.show_last_history_iters_count = 0
self.update_preview = True
def flask_thread(s2c, c2s, preview_queue): def flask_thread(s2c, c2s, s2flask, args):
config = { # config = {
"DEBUG": True, # some Flask specific configs # "DEBUG": True, # some Flask specific configs
"CACHE_TYPE": "simple", # Flask-Caching related configs # "CACHE_TYPE": "simple", # Flask-Caching related configs
"CACHE_DEFAULT_TIMEOUT": 300 # "CACHE_DEFAULT_TIMEOUT": 300
} # }
app = Flask(__name__) app = Flask(__name__)
app.config.from_mapping(config) # app.config.from_mapping(config)
cache = Cache(app) # cache = Cache(app)
template = """<html> template = """<html>
<head> <head>
<title>Video Streaming Demonstration</title> <title>Flask Server Demonstration</title>
</head> </head>
<body> <body>
<h1>Video Streaming Demonstration</h1> <h1>Video Streaming Demonstration</h1>
@ -357,8 +311,17 @@ def flask_thread(s2c, c2s, preview_queue):
</html>""" </html>"""
def gen(): def gen():
if not preview_queue.empty(): model_path = Path(args.get('model_path', ''))
frame = preview_queue.get() print('[MainThread]', 'model_path:', model_path)
filename = 'preview.jpg'
preview_file = str(model_path / filename)
print('[MainThread]', 'preview_file:', preview_file)
frame = open(preview_file, 'rb').read()
while True:
try:
frame = open(preview_file, 'rb').read()
except:
pass
yield b'--frame\r\nContent-Type: image/jpeg\r\n\r\n' yield b'--frame\r\nContent-Type: image/jpeg\r\n\r\n'
yield frame yield frame
yield b'\r\n\r\n' yield b'\r\n\r\n'
@ -371,36 +334,52 @@ def flask_thread(s2c, c2s, preview_queue):
elif 'exit' in request.form: elif 'exit' in request.form:
s2c.put({'op': 'close'}) s2c.put({'op': 'close'})
elif 'update' in request.form: elif 'update' in request.form:
while not s2flask.empty():
input = s2flask.get()
c2s.put({'op': 'update'}) c2s.put({'op': 'update'})
while s2flask.empty():
pass
input = s2flask.get()
elif 'next_preview' in request.form: elif 'next_preview' in request.form:
c2s.put({'op': 'preview'}) while not s2flask.empty():
input = s2flask.get()
c2s.put({'op': 'next_preview'})
while s2flask.empty():
pass
input = s2flask.get()
elif 'change_history_range' in request.form: elif 'change_history_range' in request.form:
while not s2flask.empty():
input = s2flask.get()
c2s.put({'op': 'change_history_range'}) c2s.put({'op': 'change_history_range'})
while s2flask.empty():
pass
input = s2flask.get()
# return '', 204
return render_template_string(template) return render_template_string(template)
def queue_not_empty():
return not preview_queue.empty()
# @app.route('/preview_image') # @app.route('/preview_image')
# @cache.cached(timeout=300, unless=queue_not_empty)
# def preview_image(): # def preview_image():
# yield Response(preview_queue.get(), # return Response(gen(), mimetype='multipart/x-mixed-replace;boundary=frame')
# mimetype='multipart/x-mixed-replace;boundary=frame')
@app.route('/preview_image') @app.route('/preview_image')
@cache.cached(timeout=300, unless=queue_not_empty)
def preview_image(): def preview_image():
return Response(preview_queue.get(), mimetype='image/jpeg') model_path = Path(args.get('model_path', ''))
filename = 'preview.jpg'
preview_file = str(model_path / filename)
return send_file(preview_file, mimetype='image/jpeg', cache_timeout=-1)
app.run(debug=True, use_reloader=False) app.run(debug=False, use_reloader=False)
def main(args, device_args): def main(args, device_args):
io.log_info ("Running trainer.\r\n") io.log_info ("Running trainer.\r\n")
no_preview = args.get('no_preview', False)
s2c = queue.Queue() s2c = queue.Queue()
c2s = queue.Queue() c2s = queue.Queue()
preview_queue = queue.Queue() s2flask = queue.Queue()
e = threading.Event() e = threading.Event()
thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e, args, device_args) ) thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e, args, device_args) )
@ -408,9 +387,82 @@ def main(args, device_args):
e.wait() #Wait for inital load to occur. e.wait() #Wait for inital load to occur.
flask_t = threading.Thread(target=flask_thread, args=(s2c, c2s, preview_queue)) flask_t = threading.Thread(target=flask_thread, args=(s2c, c2s, s2flask, args))
flask_t.start() flask_t.start()
preview = Preview(c2s, s2c, preview_queue) wnd_name = "Training preview"
preview.get_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
iteration = 0
batch_size = 1
zoom = Zoom.ZOOM_100
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
iteration = input['iter'] if 'iter' in input.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
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(args.get('model_path', ''))
filename = 'preview.jpg'
preview_file = str(model_path / filename)
cv2.imwrite(preview_file, preview_pane_image)
s2flask.put({'op': 'show'})
# socketio.emit('some event', {'data': 42})
# cv2.imshow(wnd_name, preview_pane_image)
is_showing = True
try:
io.process_messages(0.01)
except KeyboardInterrupt:
s2c.put({'op': 'close'})
io.destroy_all_windows()