evazion
069af1c0c4
Handle unexpected errors in the web service, for example if given an invalid or corrupt image file.
61 lines
1.9 KiB
Python
Executable File
61 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
from os import getenv
|
|
from dotenv import load_dotenv
|
|
from autotagger import Autotagger
|
|
from base64 import b64encode
|
|
from flask import Flask, request, render_template, jsonify, abort
|
|
from werkzeug.exceptions import HTTPException
|
|
|
|
load_dotenv()
|
|
model_path = getenv("MODEL_PATH", "models/model.pth")
|
|
autotagger = Autotagger(model_path)
|
|
|
|
app = Flask(__name__)
|
|
app.config["JSON_SORT_KEYS"] = False
|
|
app.config["JSON_PRETTYPRINT_REGULAR"] = True
|
|
|
|
@app.route("/", methods=["GET"])
|
|
def index():
|
|
return render_template("index.html")
|
|
|
|
@app.route("/evaluate", methods=["POST"])
|
|
def evaluate():
|
|
files = request.files.getlist("file")
|
|
threshold = float(request.values.get("threshold", 0.1))
|
|
output = request.values.get("format", "html")
|
|
limit = int(request.values.get("limit", 50))
|
|
|
|
predictions = autotagger.predict(files, threshold=threshold, limit=limit)
|
|
|
|
if output == "html":
|
|
for file in files:
|
|
file.seek(0)
|
|
|
|
base64data = [b64encode(file.read()).decode() for file in files]
|
|
return render_template("evaluate.html", predictions=zip(base64data, predictions))
|
|
elif output == "json":
|
|
predictions = [{ "filename": file.filename, "tags": tags } for file, tags in zip(files, predictions)]
|
|
return jsonify(predictions)
|
|
else:
|
|
abort(400)
|
|
|
|
@app.errorhandler(HTTPException)
|
|
def handle_http_exception(exception):
|
|
output = request.values.get("format", "html")
|
|
|
|
if hasattr(exception, "original_exception"):
|
|
error = exception.original_exception.__class__.__name__
|
|
message = str(exception.original_exception)
|
|
else:
|
|
error = exception.__class__.__name__
|
|
message = str(exception)
|
|
|
|
if output == "html":
|
|
return render_template("error.html", error=error, message=message)
|
|
else:
|
|
return jsonify({ "error": error, "message": message }), exception.code
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0")
|