-
-
Notifications
You must be signed in to change notification settings - Fork 16.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Allow multiple-model serving from Flask REST API (#8973)
* allow to serve multiple models using flask restapi * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup * Update restapi.py Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
- Loading branch information
1 parent
d7bc5d7
commit 4d05472
Showing
1 changed file
with
10 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,6 @@ | ||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license | ||
""" | ||
Run a Flask REST API exposing a YOLOv5s model | ||
Run a Flask REST API exposing one or more YOLOv5s models | ||
""" | ||
|
||
import argparse | ||
|
@@ -11,12 +11,13 @@ | |
from PIL import Image | ||
|
||
app = Flask(__name__) | ||
models = {} | ||
|
||
DETECTION_URL = "/v1/object-detection/yolov5s" | ||
DETECTION_URL = "/v1/object-detection/<model>" | ||
|
||
|
||
@app.route(DETECTION_URL, methods=["POST"]) | ||
def predict(): | ||
def predict(model): | ||
if request.method != "POST": | ||
return | ||
|
||
|
@@ -30,17 +31,18 @@ def predict(): | |
im_bytes = im_file.read() | ||
im = Image.open(io.BytesIO(im_bytes)) | ||
|
||
This comment has been minimized.
Sorry, something went wrong.
This comment has been minimized.
Sorry, something went wrong.
glenn-jocher
Author
Member
|
||
results = model(im, size=640) # reduce size=320 for faster inference | ||
return results.pandas().xyxy[0].to_json(orient="records") | ||
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() | ||
|
||
# Fix known issue urllib.error.HTTPError 403: rate limit exceeded https://github.com/ultralytics/yolov5/pull/7210 | ||
torch.hub._validate_not_a_forked_repo = lambda a, b, c: True | ||
for m in opt.model: | ||
models[m] = torch.hub.load("ultralytics/yolov5", m, force_reload=True, skip_validation=True) | ||
|
||
model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True) # force_reload to recache | ||
app.run(host="0.0.0.0", port=opt.port) # debug=True causes Restarting with stat |
if im.mode != 'RGB':
im = im.convert('RGB')