Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

39 lines
1.0KB

  1. """
  2. Run a rest API exposing the yolov5s object detection model
  3. """
  4. import argparse
  5. import io
  6. import torch
  7. from PIL import Image
  8. from flask import Flask, request
  9. app = Flask(__name__)
  10. DETECTION_URL = "/v1/object-detection/yolov5s"
  11. @app.route(DETECTION_URL, methods=["POST"])
  12. def predict():
  13. if not request.method == "POST":
  14. return
  15. if request.files.get("image"):
  16. image_file = request.files["image"]
  17. image_bytes = image_file.read()
  18. img = Image.open(io.BytesIO(image_bytes))
  19. results = model(img, size=640)
  20. data = results.pandas().xyxy[0].to_json(orient="records")
  21. return data
  22. if __name__ == "__main__":
  23. parser = argparse.ArgumentParser(description="Flask api exposing yolov5 model")
  24. parser.add_argument("--port", default=5000, type=int, help="port number")
  25. args = parser.parse_args()
  26. model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True).autoshape() # force_reload to recache
  27. app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat