robmarkcole glenn-jocher commited on
Commit
1479737
1 Parent(s): 1487bc8

Flask REST API Example (#2732)

Browse files

* add files

* Update README.md

* Update README.md

* Update restapi.py

pretrained=True and model.eval() are used by default when loading a model now, so no need to call them manually.

* PEP8 reformat

* PEP8 reformat

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

utils/flask_rest_api/README.md ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Flask REST API
2
+ [REST](https://en.wikipedia.org/wiki/Representational_state_transfer) [API](https://en.wikipedia.org/wiki/API)s are commonly used to expose Machine Learning (ML) models to other services. This folder contains an example REST API created using Flask to expose the `yolov5s` model from [PyTorch Hub](https://pytorch.org/hub/ultralytics_yolov5/).
3
+
4
+ ## Requirements
5
+
6
+ [Flask](https://palletsprojects.com/p/flask/) is required. Install with:
7
+ ```shell
8
+ $ pip install Flask
9
+ ```
10
+
11
+ ## Run
12
+
13
+ After Flask installation run:
14
+
15
+ ```shell
16
+ $ python3 restapi.py --port 5000
17
+ ```
18
+
19
+ Then use [curl](https://curl.se/) to perform a request:
20
+
21
+ ```shell
22
+ $ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'`
23
+ ```
24
+
25
+ The model inference results are returned:
26
+
27
+ ```shell
28
+ [{'class': 0,
29
+ 'confidence': 0.8197850585,
30
+ 'name': 'person',
31
+ 'xmax': 1159.1403808594,
32
+ 'xmin': 750.912902832,
33
+ 'ymax': 711.2583007812,
34
+ 'ymin': 44.0350036621},
35
+ {'class': 0,
36
+ 'confidence': 0.5667674541,
37
+ 'name': 'person',
38
+ 'xmax': 1065.5523681641,
39
+ 'xmin': 116.0448303223,
40
+ 'ymax': 713.8904418945,
41
+ 'ymin': 198.4603881836},
42
+ {'class': 27,
43
+ 'confidence': 0.5661227107,
44
+ 'name': 'tie',
45
+ 'xmax': 516.7975463867,
46
+ 'xmin': 416.6880187988,
47
+ 'ymax': 717.0524902344,
48
+ 'ymin': 429.2020568848}]
49
+ ```
50
+
51
+ An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py`
utils/flask_rest_api/example_request.py ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Perform test request"""
2
+ import pprint
3
+
4
+ import requests
5
+
6
+ DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
7
+ TEST_IMAGE = "zidane.jpg"
8
+
9
+ image_data = open(TEST_IMAGE, "rb").read()
10
+
11
+ response = requests.post(DETECTION_URL, files={"image": image_data}).json()
12
+
13
+ pprint.pprint(response)
utils/flask_rest_api/restapi.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Run a rest API exposing the yolov5s object detection model
3
+ """
4
+ import argparse
5
+ import io
6
+
7
+ import torch
8
+ from PIL import Image
9
+ from flask import Flask, request
10
+
11
+ app = Flask(__name__)
12
+
13
+ DETECTION_URL = "/v1/object-detection/yolov5s"
14
+
15
+
16
+ @app.route(DETECTION_URL, methods=["POST"])
17
+ def predict():
18
+ if not request.method == "POST":
19
+ return
20
+
21
+ if request.files.get("image"):
22
+ image_file = request.files["image"]
23
+ image_bytes = image_file.read()
24
+
25
+ img = Image.open(io.BytesIO(image_bytes))
26
+
27
+ results = model(img, size=640)
28
+ data = results.pandas().xyxy[0].to_json(orient="records")
29
+ return data
30
+
31
+
32
+ if __name__ == "__main__":
33
+ parser = argparse.ArgumentParser(description="Flask api exposing yolov5 model")
34
+ parser.add_argument("--port", default=5000, type=int, help="port number")
35
+ args = parser.parse_args()
36
+
37
+ model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True).autoshape() # force_reload to recache
38
+ app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat