Ви не можете вибрати більше 25 тем Теми мають розпочинатися з літери або цифри, можуть містити дефіси (-) і не повинні перевищувати 35 символів.

38 lines
1.1KB

  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) # reduce size=320 for faster inference
  20. return results.pandas().xyxy[0].to_json(orient="records")
  21. if __name__ == "__main__":
  22. parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
  23. parser.add_argument("--port", default=5000, type=int, help="port number")
  24. args = parser.parse_args()
  25. model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache
  26. app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat