123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146 |
- import base64
- import logging
- import io
- import os
- from flask import Flask, render_template, request
- from load import init_model
- from PIL import Image
- from util import decode_prob
- logger = logging.getLogger("dog_breed_classifier")
- logger.setLevel(logging.DEBUG)
- app = Flask(__name__)
- # Initialize
- MODEL_DIR = os.path.abspath("./models")
- RESNET_CONFIG = {'arch':
- os.path.join(MODEL_DIR,
- 'model.Resnet50.json'),
- 'weights':
- os.path.join(MODEL_DIR,
- 'weights.Resnet50.hdf5')}
- INCEPTION_CONFIG = {'arch':
- os.path.join(MODEL_DIR,
- 'model.inceptionv3.json'),
- 'weights':
- os.path.join(MODEL_DIR,
- 'weights.inceptionv3.h5')}
- MODELS = {'resnet': RESNET_CONFIG,
- 'inception': INCEPTION_CONFIG}
- @app.route('/index')
- @app.route('/')
- def index():
- return render_template('select_files.html')
- @app.route('/settings', methods=['GET', 'POST'])
- def settings():
- """Select Model Architecture and Initialize
- """
- global model, graph, preprocess
- # grab model selected
- model_name = request.form['model']
- config = MODELS[model_name]
- # init the model with pre-trained architecture and weights
- model, graph = init_model(config['arch'], config['weights'])
- # use the proper preprocessing method
- if model_name == 'inception':
- from util import preprocess_inception
- preprocess = preprocess_inception
- else:
- from util import preprocess_resnet
- preprocess = preprocess_resnet
- return render_template('select_files.html', model_name=model_name)
- @app.route('/predict', methods=['GET', 'POST'])
- def predict():
- """File selection and display results
- """
- if request.method == 'POST' and 'file[]' in request.files:
- global model, graph, preprocess
- # grab model selected
- model_name = request.form['model']
- config = MODELS[model_name]
- # init the model with pre-trained architecture and weights
- model, graph = init_model(config['arch'], config['weights'])
- # use the proper preprocessing method
- if model_name == 'inception':
- from util import preprocess_inception
- preprocess = preprocess_inception
- else:
- from util import preprocess_resnet
- preprocess = preprocess_resnet
- inputs = []
- files = request.files.getlist('file[]')
- for file_obj in files:
- # Check if no files uploaded
- if file_obj.filename == '':
- if len(files) == 1:
- return render_template('select_files.html')
- continue
- entry = {}
- entry.update({'filename': file_obj.filename})
- try:
- img_bytes = io.BytesIO(file_obj.stream.getvalue())
- entry.update({'data':
- Image.open(
- img_bytes
- )})
- except AttributeError:
- img_bytes = io.BytesIO(file_obj.stream.read())
- entry.update({'data':
- Image.open(
- img_bytes
- )})
- # keep image in base64 for later use
- img_b64 = base64.b64encode(img_bytes.getvalue()).decode()
- entry.update({'img': img_b64})
- inputs.append(entry)
- outputs = []
- with graph.as_default():
- for input_ in inputs:
- # convert to 4D tensor to feed into our model
- x = preprocess(input_['data'])
- # perform prediction
- out = model.predict(x)
- outputs.append(out)
- # decode output prob
- outputs = decode_prob(outputs)
- results = []
- for input_, probs in zip(inputs, outputs):
- results.append({'filename': input_['filename'],
- 'image': input_['img'],
- 'predict_probs': probs})
- return render_template('results.html', results=results)
- # if no files uploaded
- return render_template('select_files.html')
- if __name__ == '__main__':
- #app.run(debug=True)
- app.run(host="0.0.0.0")
|