support-flask-rest-api

Oneflow-Inc · Mar 7, 2023 · 43f5d2b · 43f5d2b
1 parent fd5cbeb
commit 43f5d2b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/utils/downloads.py b/utils/downloads.py
@@ -66,7 +66,7 @@ def attempt_download(file, repo="Oneflow-Inc/one-yolov5", release="v1.2.0"):
     # Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v7.0', etc.
     from utils.general import LOGGER
 
-    if file.endswith(".pt"):
+    if str(file).endswith(".pt"):
         repo, release = "ultralytics/yolov5", "v7.0"
 
     def github_assets(repository, version="latest"):

diff --git a/utils/flask_rest_api/restapi.py b/utils/flask_rest_api/restapi.py
@@ -45,6 +45,6 @@ def predict(model):
     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('Oneflow-Inc/one-yolov5','custom', 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