diff --git a/utils/flask_rest_api/README.md b/utils/flask_rest_api/README.md new file mode 100644 index 000000000000..f528cc88df5a --- /dev/null +++ b/utils/flask_rest_api/README.md @@ -0,0 +1,38 @@ +## Flask Rest API +Rest API are commonly used to expose machine learning models to other services. This folder contains an example rest API created using Flask to expose the `yolov5s` model from pytorch hub. + +Install Flask in your environment and run: + +`$ python3 restapi.py --port 5000` + +Then use [curl](https://curl.se/) to perform a request: + +`$ curl -X POST -F image=@zidane.jpg 'http://localhost:5000/v1/object-detection/yolov5s'` + +The model inference results are returned: + +``` +[{'class': 0, + 'confidence': 0.8197850585, + 'name': 'person', + 'xmax': 1159.1403808594, + 'xmin': 750.912902832, + 'ymax': 711.2583007812, + 'ymin': 44.0350036621}, + {'class': 0, + 'confidence': 0.5667674541, + 'name': 'person', + 'xmax': 1065.5523681641, + 'xmin': 116.0448303223, + 'ymax': 713.8904418945, + 'ymin': 198.4603881836}, + {'class': 27, + 'confidence': 0.5661227107, + 'name': 'tie', + 'xmax': 516.7975463867, + 'xmin': 416.6880187988, + 'ymax': 717.0524902344, + 'ymin': 429.2020568848}] +``` + +An example python script to perform inference using [requests](https://docs.python-requests.org/en/master/) is given in `example_request.py` diff --git a/utils/flask_rest_api/example_request.py b/utils/flask_rest_api/example_request.py new file mode 100644 index 000000000000..e3e4609163b7 --- /dev/null +++ b/utils/flask_rest_api/example_request.py @@ -0,0 +1,12 @@ +"""Perform test request""" +import pprint +import requests + +DETECTION_URL = "http://localhost:5000/v1/object-detection/yolov5s" +TEST_IMAGE = "zidane.jpg" + +image_data = open(TEST_IMAGE, "rb").read() + +response = requests.post(DETECTION_URL, files={"image": image_data}).json() + +pprint.pprint(response) diff --git a/utils/flask_rest_api/restapi.py b/utils/flask_rest_api/restapi.py new file mode 100644 index 000000000000..2bab7b06a5e5 --- /dev/null +++ b/utils/flask_rest_api/restapi.py @@ -0,0 +1,41 @@ +""" +Run a rest API exposing the yolov5s object detection model +""" +import argparse +import io +from PIL import Image + +import torch +from flask import Flask, request + +app = Flask(__name__) + +DETECTION_URL = "/v1/object-detection/yolov5s" + + +@app.route(DETECTION_URL, methods=["POST"]) +def predict(): + if not request.method == "POST": + return + + if request.files.get("image"): + image_file = request.files["image"] + image_bytes = image_file.read() + + img = Image.open(io.BytesIO(image_bytes)) + + results = model(img, size=640) + data = results.pandas().xyxy[0].to_json(orient="records") + return data + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Flask api exposing yolov5 model") + parser.add_argument("--port", default=5000, type=int, help="port number") + args = parser.parse_args() + + model = torch.hub.load( + "ultralytics/yolov5", "yolov5s", pretrained=True, force_reload=True + ).autoshape() # force_reload = recache latest code + model.eval() + app.run(host="0.0.0.0", port=args.port) # debug=True causes Restarting with stat