90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import logging
|
|
import tempfile
|
|
import tracemalloc
|
|
|
|
from flask import Flask, request, jsonify
|
|
from waitress import serve
|
|
|
|
from image_prediction.config import CONFIG
|
|
from image_prediction.predictor import Predictor, extract_image_metadata_pairs, classify_images
|
|
from image_prediction.response import build_response
|
|
from prometheus_client import Gauge, Counter
|
|
from prometheus_flask_exporter import PrometheusMetrics
|
|
|
|
|
|
def main():
|
|
|
|
predictor = Predictor()
|
|
logging.info("Predictor ready.")
|
|
tracemalloc.start()
|
|
|
|
app = Flask(__name__)
|
|
metrics = PrometheusMetrics(app=app, path='/prometheus')
|
|
|
|
file_counter = Counter("image_prediction_file_counter", "count processed files")
|
|
ram_metric = Gauge("image_prediction_memory_usage", "Memory usage in Mb")
|
|
|
|
@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"])
|
|
@metrics.summary('image_prediction_request_time_seconds', 'Time spent processing request')
|
|
def predict():
|
|
def do_monitoring():
|
|
file_counter.inc()
|
|
_, peak = tracemalloc.get_traced_memory()
|
|
ram_metric.set(peak / 10 ** 6)
|
|
|
|
pdf = request.data
|
|
|
|
logging.debug("Running predictor on document...")
|
|
with tempfile.NamedTemporaryFile() as tmp_file:
|
|
tmp_file.write(pdf)
|
|
image_metadata_pairs = extract_image_metadata_pairs(tmp_file.name)
|
|
try:
|
|
predictions, metadata = classify_images(predictor, image_metadata_pairs)
|
|
if CONFIG.service.monitoring_enabled:
|
|
do_monitoring()
|
|
except Exception as err:
|
|
logging.warning("Analysis failed.")
|
|
logging.exception(err)
|
|
response = jsonify("Analysis failed.")
|
|
response.status_code = 500
|
|
return response
|
|
logging.debug(f"Found images in document.")
|
|
|
|
response = jsonify(build_response(list(predictions), list(metadata)))
|
|
|
|
logging.info("Analysis completed.")
|
|
return response
|
|
|
|
run_prediction_server(app, 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)
|
|
tracemalloc.stop()
|
|
|
|
|
|
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)
|
|
|
|
main()
|