| FROM python:3.6-slim | |
| # Install deps | |
| ADD requirements.txt / | |
| RUN pip install -U -r /requirements.txt | |
| # Trigger download of resnet model from github | |
| RUN python -c "from keras.applications.resnet50 import ResNet50; ResNet50()" | |
| # Add flask app | |
| ADD main.py / | |
| CMD ["python", "./main.py"] |
| import base64 | |
| import os | |
| from io import BytesIO | |
| import flask | |
| import numpy as np | |
| from PIL import Image | |
| from keras import backend as K | |
| from keras.applications.resnet50 import ResNet50, preprocess_input, decode_predictions | |
| from keras.preprocessing.image import img_to_array | |
| app = flask.Flask(__name__) # define a predict function as an endpoint | |
| K.set_learning_phase(0) | |
| model = ResNet50() | |
| input_size = (224, 224) | |
| # prepare the PIL image. customize as you need it for the model | |
| def prepare_image(img): | |
| # convert to array | |
| img = img.resize(input_size) | |
| img = img_to_array(img) | |
| img = np.expand_dims(img, axis=0) | |
| return preprocess_input(img) | |
| def predict_images(images): | |
| # output is a list of score (each score is a list of floats) | |
| img_array = [] | |
| for image in images: | |
| image = image[image.find("base64,") + len("base64,"):] | |
| img = Image.open(BytesIO(base64.b64decode(image))) | |
| prepared_img = prepare_image(img) | |
| img_array.append(prepared_img) | |
| prepared_imgs = np.concatenate(img_array) | |
| # get the prediction | |
| preds = model.predict(prepared_imgs) | |
| return decode_predictions(preds, top=1000) # Get back all the predictions | |
| def preds_to_map(preds): | |
| output = [] | |
| for pred in preds: | |
| pred_scores = {} | |
| for (imagenetID, label, prob) in pred: | |
| pred_scores[label] = float(prob) | |
| output.append(pred_scores) | |
| return output | |
| @app.route("/predict", methods=["POST"]) | |
| def predict(): | |
| # input is an array of base64-encoded images | |
| images = flask.request.json["images"] | |
| preds = predict_images(images) | |
| return flask.jsonify(preds_to_map(preds)) # start the flask app, allow remote connections | |
| app.run(host=os.getenv('HOST', "0.0.0.0"), | |
| port=os.getenv('PORT', "1234")) |
| flask==1.1.1 | |
| tensorflow==2.0.0 | |
| Keras==2.3.1 | |
| Pillow==6.1.0 |