image-classification-service/image_prediction/flask.py

from typing import Callable

from flask import Flask, request, jsonify

from image_prediction.utils import get_logger

logger = get_logger()


def make_prediction_server(predict_fn: Callable):

    app = Flask(__name__)

    @app.route("/ready", methods=["GET"])
    def ready():
        resp = jsonify("OK")
        resp.status_code = 200
        return resp

    @app.route("/health", methods=["GET"])
    def healthy():
        resp = jsonify("OK")
        resp.status_code = 200
        return resp

    @app.route("/", methods=["POST"])
    def predict():
        pdf = request.data

        logger.debug("Running predictor on document...")
        try:
            predictions = predict_fn(pdf)
            response = jsonify(predictions)
            logger.info("Analysis completed.")
            return response
        except Exception as err:
            logger.error("Analysis failed.")
            logger.exception(err)
            response = jsonify("Analysis failed.")
            response.status_code = 500
            return response

    return app