server.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. from flask import Flask, Response, request, views
  2. from infer import inference
  3. import numpy as np
  4. import cv2
  5. app = Flask("HM-DETR")
  6. class MainView(views.MethodView):
  7. @staticmethod
  8. def get():
  9. return """
  10. <h2>DETR 在线演示</h2>
  11. <form action="" method="POST" enctype="multipart/form-data">
  12. <input type="file" name="img">
  13. <input type="submit" value="上传">
  14. </form>
  15. """
  16. @staticmethod
  17. def post():
  18. pic = request.files.get("img")
  19. if pic is None:
  20. return Response(status=414, response="key param lost")
  21. ext = pic.filename.split(".")[-1].lower()
  22. if ext not in ["jpg", "jpeg", "bmp", "png", "rgb", "tif", "tiff", "gif"]:
  23. return Response(status=413, response="file type error")
  24. content = pic.read()
  25. if len(content) > 1024 * 1024 * 5:
  26. return Response(status=410, response="image is too large")
  27. img = cv2.imdecode(np.frombuffer(content, np.uint8), 1) # noqa
  28. result = inference(img)
  29. _, result = cv2.imencode(".jpg", result) # noqa
  30. return Response(response=result.tobytes(), mimetype="image/jpg")
  31. app.add_url_rule("/", view_func=MainView.as_view("main"))
  32. if __name__ == "__main__":
  33. app.run(host="0.0.0.0", port=5080)