added updated version of serve.py

This commit is contained in:
Matthias Bisping 2022-03-31 19:38:35 +02:00
parent 268b83a1ff
commit 1581880ec6

48
src/serve.py Normal file
View File

@ -0,0 +1,48 @@
import logging
from waitress import serve
from image_prediction.config import CONFIG
from image_prediction.default_objects import load_pipeline
from image_prediction.flask import make_prediction_server
from image_prediction.utils import get_logger
from image_prediction.utils.banner import show_banner
logger = get_logger()
def main():
def predict(pdf):
# Keras service_estimator.predict stalls when service_estimator was loaded in different process;
# therefore, we re-load the model (part of the pipeline) every time we process a new document.
# https://stackoverflow.com/questions/42504669/keras-tensorflow-and-multiprocessing-in-python
logger.debug("Loading pipeline...")
pipeline = load_pipeline(verbose=CONFIG.service.verbose)
logger.debug("Running pipeline...")
return pipeline(pdf)
prediction_server = make_prediction_server(predict)
run_prediction_server(prediction_server, mode=CONFIG.webserver.mode)
def run_prediction_server(app, mode="development"):
if mode == "development":
app.run(host=CONFIG.webserver.host, port=CONFIG.webserver.port, debug=True)
elif mode == "production":
serve(app, host=CONFIG.webserver.host, port=CONFIG.webserver.port, _quiet=False)
if __name__ == "__main__":
logging_level = CONFIG.service.logging_level
logging.basicConfig(level=logging_level)
logging.getLogger("flask").setLevel(logging.ERROR)
logging.getLogger("urllib3").setLevel(logging.ERROR)
logging.getLogger("werkzeug").setLevel(logging.ERROR)
logging.getLogger("waitress").setLevel(logging.ERROR)
logging.getLogger("PIL").setLevel(logging.ERROR)
logging.getLogger("h5py").setLevel(logging.ERROR)
show_banner()
main()