from waitress import serve

from image_prediction.config import CONFIG
from image_prediction.flask import make_prediction_server
from image_prediction.pipeline import load_pipeline
from image_prediction.utils import get_logger
from image_prediction.utils.banner import show_banner


def main():
    logger = get_logger()

    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, batch_size=CONFIG.service.batch_size)
        logger.debug("Running pipeline...")
        return list(pipeline(pdf))

    prediction_server = make_prediction_server(predict)
    serve(prediction_server, host=CONFIG.webserver.host, port=CONFIG.webserver.port, _quiet=False)


if __name__ == "__main__":

    show_banner()

    main()