from flask import Flask, Response, request, views
from infer import inference
import numpy as np
import cv2

app = Flask("HM-DETR")


class MainView(views.MethodView):
    @staticmethod
    def get():
        return """
        <h2>DETR 在线演示</h2>
        <form action="" method="POST" enctype="multipart/form-data">
            <input type="file" name="img">
            <input type="submit" value="上传">
        </form>
        """

    @staticmethod
    def post():
        pic = request.files.get("img")
        if pic is None:
            return Response(status=414, response="key param lost")
        ext = pic.filename.split(".")[-1].lower()
        if ext not in ["jpg", "jpeg", "bmp", "png", "rgb", "tif", "tiff", "gif"]:
            return Response(status=413, response="file type error")
        content = pic.read()
        if len(content) > 1024 * 1024 * 5:
            return Response(status=410, response="image is too large")
        img = cv2.imdecode(np.frombuffer(content, np.uint8), 1)  # noqa
        result = inference(img)
        _, result = cv2.imencode(".jpg", result)  # noqa
        return Response(response=result.tobytes(), mimetype="image/jpg")


app.add_url_rule("/", view_func=MainView.as_view("main"))

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5080)