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

Merge pull request #70 from faceshiftlabs/experiment/flask-preview-sockets

Browse source 
Experiment/flask preview sockets
This commit is contained in:
Jeremy Hummel 2019-09-14 16:39:50 -07:00 • committed by GitHub
commit 9db0507ef8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 93 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

42
flaskr/app.py
View file

@ -1,10 +1,11 @@
from pathlib import Path
from flask import Flask, send_file, Response, render_template, render_template_string, request, g
# from flask_socketio import SocketIO
from flask_socketio import SocketIO, emit
def create_flask_app(s2c, c2s, s2flask, args):
app = Flask(__name__, template_folder="templates")
app = Flask(__name__, template_folder="templates", static_folder="static")
model_path = Path(args.get('model_path', ''))
filename = 'preview.jpg'
preview_file = str(model_path / filename)
@ -42,16 +43,23 @@ def create_flask_app(s2c, c2s, s2flask, args):
request.environ.get('werkzeug.server.shutdown')()
return '', 204
@app.route('/', methods=['GET', 'POST'])
@app.route('/update', methods=['POST'])
def update():
send(c2s, 'update')
return '', 204
@app.route('/next_preview', methods=['POST'])
def next_preview():
send(c2s, 'next_preview')
return '', 204
@app.route('/change_history_range', methods=['POST'])
def change_history_range():
send(c2s, 'change_history_range')
return '', 204
@app.route('/')
def index():
if request.method == 'POST':
if 'update' in request.form:
send_and_wait(c2s, 'update')
elif 'next_preview' in request.form:
send_and_wait(c2s, 'next_preview')
elif 'change_history_range' in request.form:
send_and_wait(c2s, 'change_history_range')
# return '', 204
return render_template('index.html')
# @app.route('/preview_image')
@ -62,7 +70,17 @@ def create_flask_app(s2c, c2s, s2flask, args):
def preview_image():
return send_file(preview_file, mimetype='image/jpeg', cache_timeout=-1)
return app
socketio = SocketIO(app)
@socketio.on('connect', namespace='/')
def test_connect():
emit('my response', {'data': 'Connected'})
@socketio.on('disconnect', namespace='/test')
def test_disconnect():
print('Client disconnected')
return socketio, app

BIN
flaskr/static/favicon.ico Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 284 KiB

75
flaskr/templates/index.html
View file

@ -1,45 +1,56 @@
<head>
<!-- <link rel="stylesheet"-->
<!-- href="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/css/bootstrap.min.css"-->
<!-- integrity="sha384-ggOyR0iXCbMQv3Xipma34MD+dH/1fQ784/j6cY/iJTQUOhcWr7x9JvoRxT2MZw1T"-->
<!-- crossorigin="anonymous">-->
<script src="https://code.jquery.com/jquery-3.4.1.min.js"
<script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.4.1/jquery.min.js"
integrity="sha256-CSXorXvZcTkaix6Yvo6HppcZGetbYMGWSFlBw8HfCJo="
crossorigin="anonymous"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/2.2.0/socket.io.js"
integrity="sha256-yr4fRk/GU1ehYJPAs8P4JlTgu0Hdsp4ZKrx8bDEDC3I="
crossorigin="anonymous"></script>
<title>Flask Server Demonstration</title>
<link rel="shortcut icon" href="{{ url_for('static', filename='favicon.ico') }}">
<script type="text/javascript">
$(function() {
$('button#save').click(function() {
$.post('/save');
});
$('button#exit').click(function() {
$.post('/exit');
});
// $('button#update').click(function() {
// $.post('/update');
// });
//
// $('button#next_preview').click(function() {
// $.post('/next_preview');
// });
//
// $('button#change_history_range').click(function() {
// $.post('/change_history_range');
// });
$(function() {
$('button#save').click(function() {
$.post("{{ url_for('save') }}");
});
$('button#exit').click(function() {
$.post("{{ url_for('exit') }}");
});
$('button#update').click(function() {
$.post("{{ url_for('update') }}");
});
$('button#next_preview').click(function() {
$.post("{{ url_for('next_preview') }}");
});
$('button#change_history_range').click(function() {
$.post("{{ url_for('change_history_range') }}");
});
const socket = io.connect('http://' + document.domain + ':' + location.port);
socket.on('preview', function(msg) {
console.log(new Date(), '- new preview -', msg);
$('img#preview').attr("src", "{{ url_for('preview_image') }}?q=" + new Date().getTime());
});
socket.on('loss', function(loss_string) {
console.log(new Date(), '- loss string -', loss_string);
$('h1#loss').html(loss_string);
});
});
</script>
</head>
<body>
<h1>Flask Server Demonstration</h1>
<form action="/" method="post">
<h1 id="loss"></h1>
<div>
<button class='btn btn-default' id='save'>Save</button>
<button class='btn btn-default' id='exit'>Exit</button>
<button class='btn btn-default' id='update' name="update" value="update">Update</button>
<button class='btn btn-default' id='next_preview' name="next_preview" value="next_preview">Next preview</button>
<button class='btn btn-default' id='change_history_range' name="change_history_range" value="change_history_range">Change History Range</button>
</form>
<img src="{{ url_for('preview_image') }}">
<button class='btn btn-default' id='update'>Update</button>
<button class='btn btn-default' id='next_preview'>Next preview</button>
<button class='btn btn-default' id='change_history_range'>Change History Range</button>
</div>
<img id='preview' src="{{ url_for('preview_image') }}">
</body>
</html>

17
main.py
View file

@ -48,8 +48,8 @@ if __name__ == "__main__":
p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")
p.set_defaults (func=process_extract)
def process_dev_extract_umd_csv(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import Extractor
@ -80,7 +80,7 @@ if __name__ == "__main__":
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.")
p.set_defaults (func=process_extract_fanseg)
"""
def process_sort(arguments):
os_utils.set_process_lowest_prio()
from mainscripts import Sorter
@ -109,7 +109,7 @@ if __name__ == "__main__":
if arguments.remove_ie_polys:
Util.remove_ie_polys_folder (input_path=arguments.input_dir)
p = subparsers.add_parser( "util", help="Utilities.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.")
@ -129,13 +129,14 @@ if __name__ == "__main__":
'model_name' : arguments.model_name,
'no_preview' : arguments.no_preview,
'debug' : arguments.debug,
'flask_preview' : arguments.flask_preview,
'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ]
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import FlaskTrainer
FlaskTrainer.main(args, device_args)
from mainscripts import FlaskTrainer as Trainer
Trainer.main(args, device_args)
p = subparsers.add_parser( "train", help="Trainer")
p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir",
@ -155,6 +156,8 @@ if __name__ == "__main__":
p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
p.add_argument('--pingpong', dest="ping_pong", default=False,
help="Cycle between a batch size of 1 and the chosen batch size")
p.add_argument('--flask-preview', action="store_true", dest="flask_preview", default=False,
help="Launches a flask server to view the previews in a web browser")
p.set_defaults (func=process_train)
def process_convert(arguments):
@ -251,7 +254,7 @@ if __name__ == "__main__":
p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.")
p.add_argument('--no-default-mask', action="store_true", dest="no_default_mask", default=False, help="Don't use default mask.")
p.set_defaults(func=process_labelingtool_edit_mask)
def bad_args(arguments):

13
mainscripts/FlaskTrainer.py
View file

@ -15,7 +15,7 @@ import models
from interact import interact as io
def trainerThread (s2c, c2s, e, args, device_args):
def trainerThread (s2c, c2s, e, args, device_args, socketio=None):
while True:
try:
start_time = time.time()
@ -138,6 +138,9 @@ def trainerThread (s2c, c2s, e, args, device_args):
else:
io.log_info (loss_string, end='\r')
if socketio is not None:
socketio.emit('loss', loss_string)
if model.get_target_iter() != 0 and model.is_reached_iter_goal():
io.log_info ('Reached target iteration.')
model_save()
@ -289,15 +292,15 @@ def main(args, device_args):
s2c = queue.Queue()
c2s = queue.Queue()
s2flask = queue.Queue()
socketio, flask_app = create_flask_app(s2c, c2s, s2flask, args)
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, socketio))
thread.start()
e.wait() #Wait for inital load to occur.
flask_app = create_flask_app(s2c, c2s, s2flask, args)
flask_t = threading.Thread(target=flask_app.run, kwargs={'debug': True, 'use_reloader': False})
flask_t = threading.Thread(target=socketio.run, args=(flask_app,), kwargs={'debug': True, 'use_reloader': False})
flask_t.start()
wnd_name = "Training preview"
@ -366,7 +369,7 @@ def main(args, device_args):
preview_file = str(model_path / filename)
cv2.imwrite(preview_file, preview_pane_image)
s2flask.put({'op': 'show'})
# socketio.emit('some event', {'data': 42})
socketio.emit('preview', {'iter': iteration, 'loss': loss_history[-1]})
# cv2.imshow(wnd_name, preview_pane_image)
is_showing = True

1
mainscripts/Trainer.py
View file

@ -285,6 +285,7 @@ def main(args, device_args):
io.log_info ("Running trainer.\r\n")
no_preview = args.get('no_preview', False)
flask_preview = args.get('flask_preview', False)
s2c = queue.Queue()

1
requirements-cpu.txt
View file

@ -8,3 +8,4 @@ tqdm
ffmpeg-python==0.1.17
git+https://www.github.com/keras-team/keras-contrib.git
Flask==1.1.1
flask-socketio==4.2.1