YOLOv5 部署到 Web 端(Flask + JavaScript 简单易懂)

YOLOv5 是一个高效的目标检测模型,广泛应用于计算机视觉任务。本文将介绍如何使用 Flask 框架将 YOLOv5 部署到 Web 端,并通过简单的 JavaScript 前端展示检测结果。

环境准备

首先,你需要安装以下库和工具:

  1. Python 3.x
  2. Flask
  3. PyTorch
  4. 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 前端实现图像的目标检测。该示例提供了基础的结构,可以根据实际需求进一步扩展功能,例如增加日志、异常处理以及前端界面的更多优化。希望这能为你在计算机视觉领域的项目提供一些帮助。

点赞(0) 打赏

微信小程序

微信扫一扫体验

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部