YOLOv5 部署到 Web 端(Flask + JavaScript 简单易懂)
YOLOv5 是一个高效的目标检测模型,广泛应用于计算机视觉任务。本文将介绍如何使用 Flask 框架将 YOLOv5 部署到 Web 端,并通过简单的 JavaScript 前端展示检测结果。
环境准备
首先,你需要安装以下库和工具:
- Python 3.x
- Flask
- PyTorch
- YOLOv5 代码库
可以使用以下命令安装 Flask 和 PyTorch:
pip install flask
pip install torch torchvision torchaudio
然后,从 YOLOv5 的 GitHub 仓库克隆代码:
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt
Flask 应用程序
接下来,我们将创建一个简单的 Flask 应用。这个应用会接收图像文件,使用 YOLOv5 进行目标检测,并返回带有检测结果的图像。
1. 创建 Flask 应用
在 yolov5
文件夹中创建一个 app.py
文件:
from flask import Flask, request, jsonify
import torch
from PIL import Image
import io
app = Flask(__name__)
# 加载 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # 选择合适的模型
@app.route('/predict', methods=['POST'])
def predict():
if 'file' not in request.files:
return jsonify({'error': '没有文件上传'})
file = request.files['file']
img_bytes = file.read()
# 使用 YOLOv5 进行推理
img = Image.open(io.BytesIO(img_bytes))
results = model(img)
# 将结果转换为 JSON 格式
result_json = results.pandas().xyxy[0].to_dict(orient="records")
return jsonify(result_json)
if __name__ == "__main__":
app.run(debug=True)
2. 启动 Flask 应用
在终端中运行 Flask 应用:
python app.py
默认情况下,Flask 将在 http://127.0.0.1:5000/
启动。
前端 HTML 页面
现在,我们需要创建一个简单的前端 HTML 页面,让用户可以上传图像并查看检测结果。
在 yolov5
文件夹中创建一个 templates
文件夹,并在其中创建 index.html
文件:
<!DOCTYPE html>
<html lang="zh">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>YOLOv5目标检测</title>
</head>
<body>
<h1>YOLOv5 目标检测</h1>
<form id="upload-form">
<input type="file" id="file-input" name="file" accept="image/*" required>
<button type="submit">上传并检测</button>
</form>
<h2>检测结果</h2>
<pre id="result"></pre>
<script>
document.getElementById('upload-form').onsubmit = async function(event) {
event.preventDefault();
const formData = new FormData();
const fileInput = document.getElementById('file-input');
formData.append('file', fileInput.files[0]);
const response = await fetch('/predict', {
method: 'POST',
body: formData
});
const result = await response.json();
document.getElementById('result').textContent = JSON.stringify(result, null, 2);
};
</script>
</body>
</html>
3. 修改 Flask 应用以渲染 HTML
在 app.py
中添加一个路由,以便能够访问 HTML 页面:
from flask import render_template
@app.route('/')
def index():
return render_template('index.html')
运行应用
再次确保 Flask 应用正在运行,然后在浏览器中访问 http://127.0.0.1:5000/
。你将看到图像上传表单,上传图像后,检测结果将以 JSON 格式显示在页面上。
总结
本文展示了如何将 YOLOv5 部署到 Web 端,通过 Flask 和简单的 JavaScript 前端实现图像的目标检测。该示例提供了基础的结构,可以根据实际需求进一步扩展功能,例如增加日志、异常处理以及前端界面的更多优化。希望这能为你在计算机视觉领域的项目提供一些帮助。