Sa-m commited on
Commit
85b1295
·
1 Parent(s): 74e38ac

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +32 -0
app.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Run a rest API exposing the yolov5s object detection model
3
+ """
4
+ import argparse
5
+ import io
6
+ import torch
7
+ from flask import Flask, request
8
+ from PIL import Image
9
+
10
+ app = Flask(__name__)
11
+
12
+ DETECTION_URL = "/v1/detect"
13
+
14
+
15
+ @app.route(DETECTION_URL, methods=["POST"])
16
+ def predict():
17
+ if not request.method == "POST":
18
+ return
19
+
20
+ if request.files.get("image"):
21
+ image_file = request.files["image"]
22
+ image_bytes = image_file.read()
23
+
24
+ img = Image.open(io.BytesIO(image_bytes))
25
+
26
+ results = model(img, size=640) # reduce size=320 for faster inference
27
+ return results.pandas().xyxy[0].to_json(orient="records")
28
+
29
+
30
+ if __name__ == "__main__":
31
+ model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True) # force_reload to recache
32
+ app.run(host="0.0.0.0", port=7860,debug =True) # debug=True causes Restarting with stat