-
Notifications
You must be signed in to change notification settings - Fork 3
/
app.py
62 lines (49 loc) · 2.09 KB
/
app.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
import os
from flask import Flask, flash, request, redirect, url_for, render_template
from datetime import datetime
# Constants
UPLOAD_FOLDER = './static'
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'jpeg'])
MODEL = 'densenet'
app = Flask(__name__)
app.secret_key = b'_5ef8t202w@$Trz\n\xec]/'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
def allowed_file(filename):
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@app.route('/', methods=['GET', 'POST'])
def upload_file():
if request.method == 'POST':
if 'file' not in request.files:
flash('No file part')
return redirect(request.url)
file = request.files['file']
if file.filename == '':
flash('No selected file')
return redirect(request.url)
if file and allowed_file(file.filename):
# Save the input image
filename = "input.png"
file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
# Input and Output filenames
input_filename = './static/input.png'
models = request.form.getlist('model')
if models == []:
flash('No ML model selected')
return redirect(request.url)
MODELS = ['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'densenet']
model_dict = {}
for MODEL in MODELS:
if MODEL in models:
model_dict[MODEL] = True
# Run inference.py
script_command = f'python3 inference.py -d {input_filename} -o ./static/output_{MODEL}.png -be {MODEL}'
os.system(script_command)
else:
model_dict[MODEL] = False
# Add dummy values to image path to avoid HTML Image caching
suffix = "?" + datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
return render_template('display.html', suffix = suffix, model_dict = model_dict)
return render_template('home.html')
if __name__ == "__main__":
app.run(debug=True)