Skip to content

Commit

Permalink
support-flask-rest-api in utils/flask_rest_api/restapi.py
Browse files Browse the repository at this point in the history
  • Loading branch information
ccssu committed Mar 7, 2023
1 parent c0fb352 commit fd5cbeb
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 4 deletions.
2 changes: 1 addition & 1 deletion utils/flask_rest_api/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

[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/).
created using Flask to expose the YOLOv5s model from [OneFlow Hub](https://oneflow.readthedocs.io/en/master/hub.html).

## Requirements

Expand Down
8 changes: 5 additions & 3 deletions utils/flask_rest_api/restapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,14 @@

import argparse
import io

import ssl
import oneflow as torch
from flask import Flask, request
from PIL import Image

# https://stackoverflow.com/questions/50236117/scraping-ssl-certificate-verify-failed-error-for-http-en-wikipedia-org
ssl._create_default_https_context = ssl._create_unverified_context

app = Flask(__name__)
models = {}

Expand Down Expand Up @@ -41,8 +44,7 @@ def predict(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)
models[m] = torch.hub.load('Oneflow-Inc/one-yolov5','custom', m ,force_reload=True, skip_validation=True)

app.run(host="0.0.0.0", port=opt.port) # debug=True causes Restarting with stat

0 comments on commit fd5cbeb

Please sign in to comment.