├── static ├── pytorch.png └── style.css ├── README.md ├── LICENSE ├── app.py └── templates ├── index.html └── result.html /static/pytorch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jzhang533/yolov5-flask/HEAD/static/pytorch.png -------------------------------------------------------------------------------- /static/style.css: -------------------------------------------------------------------------------- 1 | html, 2 | body { 3 | height: 100%; 4 | } 5 | 6 | body { 7 | display: -ms-flexbox; 8 | display: flex; 9 | -ms-flex-align: center; 10 | align-items: center; 11 | padding-top: 40px; 12 | padding-bottom: 40px; 13 | background-color: #f5f5f5; 14 | } 15 | 16 | .form-signin { 17 | width: 100%; 18 | max-width: 330px; 19 | padding: 15px; 20 | margin: auto; 21 | } 22 | 23 | .form-signin .form-control { 24 | position: relative; 25 | box-sizing: border-box; 26 | height: auto; 27 | padding: 10px; 28 | font-size: 16px; 29 | } 30 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # deploy a yolov5 model using pytorch + flask 2 | 3 | ## start 4 | 5 | ```bash 6 | $ FLASK_ENV=development FLASK_APP=app.py flask run 7 | ``` 8 | then, visit http://localhost:5000/ in your browser 9 | 10 | ## environments 11 | 12 | - pytorch >= 1.6 13 | - flask 14 | - and [dependencies](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) required by [ultralytics/yolov5](https://github.com/ultralytics/yolov5) 15 | 16 | ## reference 17 | - https://github.com/ultralytics/yolov5 18 | - [Load YOLOv5 from PyTorch Hub ](https://github.com/ultralytics/yolov5/issues/36) 19 | - https://github.com/avinassh/pytorch-flask-api-heroku 20 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 jzhang533 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /app.py: -------------------------------------------------------------------------------- 1 | import io 2 | import os 3 | import json 4 | from PIL import Image 5 | 6 | import torch 7 | from flask import Flask, jsonify, url_for, render_template, request, redirect 8 | 9 | app = Flask(__name__) 10 | 11 | RESULT_FOLDER = os.path.join('static') 12 | app.config['RESULT_FOLDER'] = RESULT_FOLDER 13 | 14 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) 15 | model.eval() 16 | 17 | def get_prediction(img_bytes): 18 | img = Image.open(io.BytesIO(img_bytes)) 19 | imgs = [img] # batched list of images 20 | 21 | # Inference 22 | results = model(imgs, size=640) # includes NMS 23 | return results 24 | 25 | @app.route('/', methods=['GET', 'POST']) 26 | def predict(): 27 | if request.method == 'POST': 28 | if 'file' not in request.files: 29 | return redirect(request.url) 30 | file = request.files.get('file') 31 | if not file: 32 | return 33 | 34 | img_bytes = file.read() 35 | results = get_prediction(img_bytes) 36 | 37 | results.save(save_dir='static') 38 | 39 | full_filename = os.path.join(app.config['RESULT_FOLDER'], 'results0.jpg') 40 | return redirect('static/image0.jpg') 41 | return render_template('index.html') 42 | -------------------------------------------------------------------------------- /templates/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 19 | 20 | 21 | yolov5 object deteection using PyTorch 22 | 23 | 24 |
25 | 26 |

Upload any image

27 | 28 |
29 | 30 |

Built using Pytorch, Flask and Love

31 |
32 | 33 | 34 | 35 | 51 | 52 | 53 | -------------------------------------------------------------------------------- /templates/result.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 19 | 20 | 21 | Image Prediction using PyTorch 22 | 23 | 24 | User Image 25 |
26 | 27 |

Prediction

28 |
Detected Image: {{ class_name }}
29 |
ImageNet Class ID: {{ class_id }}
30 |

Built using Pytorch, Flask and Love

31 |
32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | --------------------------------------------------------------------------------