resolve issue #1

jzhang533 · Jul 13, 2022 · 1713567 · 1713567
1 parent fb347c9
commit 1713567
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -11,7 +11,7 @@ then, visit http://localhost:5000/ in your browser
 
 - pytorch >= 1.6
 - flask
-- pillow
+- and [dependencies](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) required by [ultralytics/yolov5](https://github.com/ultralytics/yolov5)
 
 ## reference
 - https://github.com/ultralytics/yolov5

diff --git a/app.py b/app.py
@@ -11,7 +11,7 @@
 RESULT_FOLDER = os.path.join('static')
 app.config['RESULT_FOLDER'] = RESULT_FOLDER
 
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).autoshape()  # for PIL/cv2/np inputs and NMS
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 model.eval()
 
 def get_prediction(img_bytes):
@@ -33,9 +33,9 @@ def predict():
 
         img_bytes = file.read()
         results = get_prediction(img_bytes)
-        results.save()  # save as results1.jpg, results2.jpg... etc.
-        os.rename("results0.jpg", "static/results0.jpg")
+
+        results.save(save_dir='static')
 
         full_filename = os.path.join(app.config['RESULT_FOLDER'], 'results0.jpg')
-        return redirect('static/results0.jpg')
+        return redirect('static/image0.jpg')
     return render_template('index.html')