first commit
This commit is contained in:
73
app/yolov5/utils/flask_rest_api/README.md
Normal file
73
app/yolov5/utils/flask_rest_api/README.md
Normal file
@ -0,0 +1,73 @@
|
||||
# Flask REST API
|
||||
|
||||
[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/).
|
||||
|
||||
## Requirements
|
||||
|
||||
[Flask](https://palletsprojects.com/p/flask/) is required. Install with:
|
||||
|
||||
```shell
|
||||
$ pip install Flask
|
||||
```
|
||||
|
||||
## Run
|
||||
|
||||
After Flask installation run:
|
||||
|
||||
```shell
|
||||
$ python3 restapi.py --port 5000
|
||||
```
|
||||
|
||||
Then use [curl](https://curl.se/) to perform a request:
|
||||
|
||||
```shell
|
||||
$ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'
|
||||
```
|
||||
|
||||
The model inference results are returned as a JSON response:
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"class": 0,
|
||||
"confidence": 0.8900438547,
|
||||
"height": 0.9318675399,
|
||||
"name": "person",
|
||||
"width": 0.3264600933,
|
||||
"xcenter": 0.7438579798,
|
||||
"ycenter": 0.5207948685
|
||||
},
|
||||
{
|
||||
"class": 0,
|
||||
"confidence": 0.8440024257,
|
||||
"height": 0.7155083418,
|
||||
"name": "person",
|
||||
"width": 0.6546785235,
|
||||
"xcenter": 0.427829951,
|
||||
"ycenter": 0.6334488392
|
||||
},
|
||||
{
|
||||
"class": 27,
|
||||
"confidence": 0.3771208823,
|
||||
"height": 0.3902671337,
|
||||
"name": "tie",
|
||||
"width": 0.0696444362,
|
||||
"xcenter": 0.3675483763,
|
||||
"ycenter": 0.7991207838
|
||||
},
|
||||
{
|
||||
"class": 27,
|
||||
"confidence": 0.3527112305,
|
||||
"height": 0.1540903747,
|
||||
"name": "tie",
|
||||
"width": 0.0336618312,
|
||||
"xcenter": 0.7814827561,
|
||||
"ycenter": 0.5065554976
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given
|
||||
in `example_request.py`
|
19
app/yolov5/utils/flask_rest_api/example_request.py
Normal file
19
app/yolov5/utils/flask_rest_api/example_request.py
Normal file
@ -0,0 +1,19 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Perform test request
|
||||
"""
|
||||
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s"
|
||||
IMAGE = "zidane.jpg"
|
||||
|
||||
# Read image
|
||||
with open(IMAGE, "rb") as f:
|
||||
image_data = f.read()
|
||||
|
||||
response = requests.post(DETECTION_URL, files={"image": image_data}).json()
|
||||
|
||||
pprint.pprint(response)
|
48
app/yolov5/utils/flask_rest_api/restapi.py
Normal file
48
app/yolov5/utils/flask_rest_api/restapi.py
Normal file
@ -0,0 +1,48 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Run a Flask REST API exposing one or more YOLOv5s models
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
|
||||
import torch
|
||||
from flask import Flask, request
|
||||
from PIL import Image
|
||||
|
||||
app = Flask(__name__)
|
||||
models = {}
|
||||
|
||||
DETECTION_URL = "/v1/object-detection/<model>"
|
||||
|
||||
|
||||
@app.route(DETECTION_URL, methods=["POST"])
|
||||
def predict(model):
|
||||
if request.method != "POST":
|
||||
return
|
||||
|
||||
if request.files.get("image"):
|
||||
# Method 1
|
||||
# with request.files["image"] as f:
|
||||
# im = Image.open(io.BytesIO(f.read()))
|
||||
|
||||
# Method 2
|
||||
im_file = request.files["image"]
|
||||
im_bytes = im_file.read()
|
||||
im = Image.open(io.BytesIO(im_bytes))
|
||||
|
||||
if model in models:
|
||||
results = models[model](im, size=640) # reduce size=320 for faster inference
|
||||
return results.pandas().xyxy[0].to_json(orient="records")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
|
||||
parser.add_argument("--port", default=5000, type=int, help="port number")
|
||||
parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')
|
||||
opt = parser.parse_args()
|
||||
|
||||
for m in opt.model:
|
||||
models[m] = torch.hub.load("ultralytics/yolov5", m, force_reload=True, skip_validation=True)
|
||||
|
||||
app.run(host="0.0.0.0", port=opt.port) # debug=True causes Restarting with stat
|
Reference in New Issue
Block a user