devbox
秋刀鱼在做梦
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

yolov5使用flask部署至前端，实现照片\视频识别_yolov5与flask集成

作者：秋刀鱼在做梦 | 2024-07-11 02:50:17

赞

踩

yolov5与flask集成
大半年前初学yolo flask时，急需此功能，Csdn、Github、B站找到很多教程，效果并不是很满意。
近期做项目碰到类似需求，再度尝试，实现简单功能，分享下相关代码，仅学习使用，如有纰漏，望多包涵。

实现功能：

  1. 可更换权重文件（best.py）

  2. 上传图片并识别,可以点击图片放大查看

  3. 上传视频并识别

  4. 识别后的文件下载功能

效果图如上

文件结构如下：

project/
  static/

  空

  templates/
    index.html
    
 app.py
 

相关代码：

app.py

  1. import cv2
  2. import numpy as np
  3. import torch
  4. from flask import Flask, request, jsonify, render_template
  5. import base64
  6. import os
  7. from datetime import datetime
  8. from werkzeug.utils import secure_filename
  9. app = Flask(__name__)
  10.  
  11. # 全局变量：模型和模型权重路径
  12. model = None
  13.  
  14.  
  15.  
  16. # 提前加载模型
  17. # 提前加载模型
  18. def load_model():
  19.     global model
  20.     global new_filename
  21.     # 拼接权重文件的完整路径
  22.     model = torch.hub.load("E:\\pythonProject2\\flaskProject\\yolov5-master", "custom", path='weight/'+new_filename, source="local")
  23.  
  24.  
  25. # 路由处理图片检测请求
  26. @app.route("/predict_image", methods=["POST"])
  27. def predict_image():
  28.     global model
  29.  
  30.     # 获取图像文件
  31.     file = request.files["image"]
  32.     # 读取图像数据并转换为RGB格式
  33.     image_data = file.read()
  34.     nparr = np.frombuffer(image_data, np.uint8)
  35.     image = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
  36.  
  37.     results = model(image)
  38.  
  39.     image = results.render()[0]
  40.  
  41.     # 将图像转换为 base64 编码的字符串
  42.     _, buffer = cv2.imencode(".png", image)
  43.     image_str = base64.b64encode(buffer).decode("utf-8")
  44.  
  45.     # 获取当前时间，并将其格式化为字符串
  46.     current_time = datetime.now().strftime("%Y%m%d%H%M%S")
  47.     # 构建保存路径
  48.     save_dir = "static"
  49.     filename, extension = os.path.splitext(file.filename)  # 获取上传文件的文件名和扩展名
  50.     save_filename = f"{filename}_{current_time}{extension}"
  51.     save_path = os.path.join(save_dir, save_filename)
  52.     cv2.imwrite(save_path, image)
  53.  
  54.     return jsonify({"image": image_str})
  55.  
  56.  
  57. # 函数用于在视频帧上绘制检测结果
  58. def detect_objects(frame, model):
  59.     results = model(frame)
  60.     detections = results.pred[0]  # 这里假设只有一张输入图片
  61.     # 在帧上绘制检测结果
  62.     for det in detections:
  63.         # 获取边界框信息
  64.         x1, y1, x2, y2, conf, class_id = det[:6]
  65.  
  66.         # 在帧上绘制边界框
  67.         cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
  68.  
  69.         # 在帧上绘制类别和置信度
  70.         label = f'{model.names[int(class_id)]} {conf:.2f}'
  71.         cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
  72.         print(f'Confidence: {conf:.2f}')
  73.     return frame
  74.  
  75.  
  76. # 路由处理视频检测请求
  77. @app.route("/predict_video", methods=["POST"])
  78. def predict_video():
  79.     global model
  80.  
  81.     # 从请求中获取视频文件
  82.     video_file = request.files["video"]
  83.     # 保存视频到临时文件
  84.     temp_video_path = "temp_video.mp4"
  85.     video_file.save(temp_video_path)
  86.  
  87.     # 逐帧读取视频
  88.     video = cv2.VideoCapture(temp_video_path)
  89.  
  90.     # 获取视频的帧率和尺寸
  91.     fps = video.get(cv2.CAP_PROP_FPS)
  92.     width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
  93.     height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
  94.  
  95.     # 视频写入对象
  96.     output_video_filename = f"output_video_{datetime.now().strftime('%Y%m%d%H%M%S')}.mp4"
  97.     output_video_path = os.path.join("static", output_video_filename)
  98.     fourcc = cv2.VideoWriter_fourcc(*"avc1")  # 使用 H.264 编码器
  99.     out_video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
  100.  
  101.     # 逐帧处理视频并进行目标检测
  102.     while True:
  103.         ret, frame = video.read()
  104.         if not ret:
  105.             break
  106.  
  107.         # 进行目标检测
  108.         detection_result = detect_objects(frame, model)
  109.  
  110.         # 将处理后的帧写入输出视频
  111.         out_video.write(detection_result)
  112.  
  113.     # 释放视频对象
  114.     video.release()
  115.     out_video.release()
  116.  
  117.     return jsonify({"output_video_path": output_video_filename})
  118.  
  119. @app.route("/upload_weight", methods=["POST"])
  120. def upload_weight():
  121.     global new_filename
  122.     # 获取上传的权重文件
  123.     weight_file = request.files["weight"]
  124.     # 获取上传文件的原始文件名
  125.     original_filename = secure_filename(weight_file.filename)
  126.     # 提取文件名和扩展名
  127.     filename, extension = os.path.splitext(original_filename)
  128.     # 构造新的文件名，加上当前时间戳
  129.     current_time = datetime.now().strftime("%Y%m%d%H%M%S")
  130.     new_filename = f"best_{current_time}.pt"
  131.     # 拼接权重文件的保存路径
  132.     save_path = os.path.join("E:\\pythonProject2\\flaskProject\\weight\\", new_filename)
  133.     # 保存权重文件
  134.     weight_file.save(save_path)
  135.     # 加载模型
  136.     load_model()
  137.  
  138.     return jsonify({"message": "Weight file uploaded successfully and model loaded"})
  139. @app.route("/")
  140. def index():
  141.  
  142.     return render_template("index.html")
  143.  
  144. if __name__ == "__main__":
  145.     app.run(debug=True)
  146.

index.html

  1. <!DOCTYPE html>
  2. <html lang="en">
  3. <head>
  4.     <meta charset="UTF-8">
  5.     <meta name="viewport" content="width=device-width, initial-scale=1.0">
  6.     <title>Object Detection</title>
  7.     <style>
  8.         body {
  9.             font-family: Arial, sans-serif;
  10.             margin: 0;
  11.             padding: 0;
  12.             background-color: #f3f3f3;
  13.             display: flex;
  14.             justify-content: center;
  15.             align-items: center;
  16.             height: 100vh;
  17.             flex-direction: column;
  18.         }
  19.  
  20.         #content {
  21.             text-align: center;
  22.             max-width: 820px;
  23.             margin-top: 20px;
  24.         }
  25.  
  26.         h1 {
  27.             color: #333;
  28.         }
  29.  
  30.         h2 {
  31.             color: #666;
  32.         }
  33.  
  34.         input[type="file"] {
  35.             margin-bottom: 10px;
  36.         }
  37.  
  38.         .media-container {
  39.             display: flex;
  40.             max-width: 100%;
  41.             margin-bottom: 20px;
  42.         }
  43.  
  44.         .media-container:first-child {
  45.             margin-right: 20px; /* 在第一个容器的右侧添加间隔 */
  46.         }
  47.  
  48.         .media-container img,
  49.         .media-container video {
  50.             max-width: 100%;
  51.             height: auto;
  52.         }
  53.  
  54.         .original {
  55.             width: 400px;
  56.             overflow: hidden;
  57.         }
  58.  
  59.         .processed {
  60.             flex: 2; /* 右边容器占据剩余空间 */
  61.         }
  62.  
  63.         button {
  64.             padding: 10px 20px;
  65.             background-color: #007bff;
  66.             color: #fff;
  67.             border: none;
  68.             border-radius: 5px;
  69.             cursor: pointer;
  70.             margin-bottom: 10px;
  71.         }
  72.  
  73.         /* 新增样式：模态框 */
  74.         .modal {
  75.             display: none; /* 默认隐藏 */
  76.             position: fixed;
  77.             z-index: 1;
  78.             left: 0;
  79.             top: 0;
  80.             width: 100%;
  81.             height: 100%;
  82.             overflow: auto;
  83.             background-color: rgba(0, 0, 0, 0.9); /* 半透明黑色背景 */
  84.         }
  85.  
  86.         .modal-content {
  87.             margin: auto;
  88.             display: block;
  89.             width: 80%;
  90.             max-width: 800px;
  91.             position: absolute;
  92.             left: 50%;
  93.             top: 50%;
  94.             transform: translate(-50%, -50%);
  95.             text-align: center; /* 居中显示图片 */
  96.         }
  97.  
  98.         .close {
  99.             color: #ccc;
  100.             font-size: 36px;
  101.             font-weight: bold;
  102.             cursor: pointer;
  103.             position: absolute;
  104.             top: 10px;
  105.             right: 10px;
  106.         }
  107.  
  108.         .close:hover,
  109.         .close:focus {
  110.             color: #fff;
  111.             text-decoration: none;
  112.         }
  113.         #downloadButton {
  114.            padding: 10px 20px;
  115.             background-color: #007bff;
  116.             color: #fff;
  117.             border: none;
  118.             border-radius: 5px;
  119.             cursor: pointer;
  120.             margin-bottom: 10px;
  121.  
  122.         }
  123.  
  124.         /* 新增样式：响应式图片 */
  125.         .modal-content img,
  126.         .modal-content video {
  127.             max-width: 100%;
  128.             height: auto;
  129.         }
  130.  
  131.     </style>
  132. </head>
  133. <body>
  134.     <h2>上传权重文件</h2>
  135.         <!-- 新增按钮用于触发上传权重文件 -->
  136.         <button onclick="document.getElementById('weightFile').click()">选择权重文件</button>
  137.         <input type="file" id="weightFile" accept=".pt" onchange="displaySelectedWeightFile()" style="display: none;">
  138.         <br>
  139.     <!-- 新增模态框 -->
  140.     <div id="myModal" class="modal" onclick="closeModal()">
  141.         <div class="modal-content" id="modalContent" onclick="stopPropagation(event)">
  142.             <!-- 放大后的图片或视频将在这里显示 -->
  143.             <span class="close" onclick="closeModal()">&times;</span>
  144.         </div>
  145.     </div>
  146.  
  147.     <div id="content">
  148.         <h1>照片/视频检测</h1>
  149.  
  150.         <!-- 上传图片 -->
  151.         <h2>上传图片</h2>
  152.         <input type="file" id="imageFile" accept="image/*" onchange="displaySelectedImage()">
  153.         <button onclick="uploadImage()">上传</button>
  154.         <button id="downloadImageButton"  onclick="downloadProcessedImage()">下载</button>
  155.         <br>
  156.         <div class="media-container">
  157.             <div class="original media-container" onclick="enlargeImage()">
  158.                 <img id="uploadedImage" src="#" alt="Uploaded Image" style="display:none;">
  159.                 <button id="zoomInButton" style="display:none;">Zoom In</button>
  160.             </div>
  161.             <div class="processed media-container" onclick="enlargeImage2()">
  162.                 <img id="processedImage" src="#" alt="Processed Image" style="display:none;">
  163.  
  164.             </div>
  165.         </div>
  166.         <br>
  167.  
  168.         <!-- 上传视频 -->
  169.         <h2>上传视频</h2>
  170.         <input type="file" id="videoFile" accept="video/mp4,video/x-m4v,video/*" onchange="displaySelectedVideo()">
  171.         <button onclick="uploadVideo()">上传</button>
  172.         <button id="downloadButton" onclick="downloadProcessedVideo()">下载</button>
  173.         <br>
  174.         <div class="media-container">
  175.             <div class="original media-container" >
  176.                 <video id="uploadedVideo" src="#" controls style="display:none;"></video>
  177.             </div>
  178.             <div class="processed media-container">
  179.                 <video id="processedVideo" controls style="display:none;"></video>
  180.  
  181.             </div>
  182.         </div>
  183.         <br>
  184.  
  185.     </div>
  186.  
  187.     <script>
  188.          // 显示选择的权重文件
  189.         function displaySelectedWeightFile() {
  190.             var fileInput = document.getElementById('weightFile');
  191.             var file = fileInput.files[0];
  192.             console.log('Selected weight file:', file);
  193.             // 上传权重文件
  194.             uploadWeight(file);
  195.         }
  196.  
  197.         // 上传权重文件
  198.         function uploadWeight(file) {
  199.             var formData = new FormData();
  200.             formData.append('weight', file);
  201.  
  202.             fetch('/upload_weight', {
  203.                 method: 'POST',
  204.                 body: formData
  205.             })
  206.             .then(response => response.json())
  207.             .then(data => {
  208.                 console.log('Upload weight response:', data);
  209.                 // 可以根据后端返回的响应进行相应的处理
  210.             })
  211.             .catch(error => console.error('Error:', error));
  212.         }
  213.         // 显示选择的图片并添加点击放大功能
  214.         function displaySelectedImage() {
  215.             var fileInput = document.getElementById('imageFile');
  216.             var file = fileInput.files[0];
  217.             var imageElement = document.getElementById('uploadedImage');
  218.             imageElement.src = URL.createObjectURL(file);
  219.             imageElement.style.display = 'inline';
  220.             document.getElementById('zoomInButton').style.display = 'inline';
  221.         }
  222.  
  223.         // 显示模态框并放大图片
  224.         function enlargeImage() {
  225.             var modal = document.getElementById('myModal');
  226.             var modalImg = document.getElementById('modalContent');
  227.             var img = document.getElementById('uploadedImage');
  228.             modal.style.display = 'block';
  229.             modalImg.innerHTML = '<img src="' + img.src + '">';
  230.         }
  231.         // 显示模态框并放大图片
  232.         function enlargeImage2() {
  233.             var modal = document.getElementById('myModal');
  234.             var modalImg = document.getElementById('modalContent');
  235.             var img = document.getElementById('processedImage');
  236.             modal.style.display = 'block';
  237.             modalImg.innerHTML = '<img src="' + img.src + '">';
  238.         }
  239.  
  240.  
  241.         // 显示选择的视频并添加点击放大功能
  242.         function displaySelectedVideo() {
  243.             var fileInput = document.getElementById('videoFile');
  244.             var file = fileInput.files[0];
  245.             var videoElement = document.getElementById('uploadedVideo');
  246.             videoElement.src = URL.createObjectURL(file);
  247.             videoElement.style.display = 'block';
  248.         }
  249.  
  250.  
  251.         // 上传图片并向后端发送请求
  252.         function uploadImage() {
  253.             var fileInput = document.getElementById('imageFile');
  254.             var file = fileInput.files[0];
  255.             var formData = new FormData();
  256.             formData.append('image', file);
  257.  
  258.             fetch('/predict_image', {
  259.                 method: 'POST',
  260.                 body: formData
  261.             })
  262.             .then(response => response.json())
  263.             .then(data => {
  264.                 var imageElement = document.getElementById('processedImage');
  265.                 imageElement.src = 'data:image/png;base64,' + data.image;
  266.                 imageElement.style.display = 'inline';
  267.                 document.getElementById('downloadImageButton').style.display = 'inline';
  268.             })
  269.             .catch(error => console.error('Error:', error));
  270.         }
  271.  
  272.         // 下载处理后的图片
  273.         function downloadProcessedImage() {
  274.             var imageElement = document.getElementById('processedImage');
  275.             var url = imageElement.src;
  276.             var a = document.createElement('a');
  277.             a.href = url;
  278.             a.download = 'processed_image.png';
  279.             document.body.appendChild(a);
  280.             a.click();
  281.             document.body.removeChild(a);
  282.         }
  283.  
  284.         // 上传视频并向后端发送请求
  285.         function uploadVideo() {
  286.             var fileInput = document.getElementById('videoFile');
  287.             var file = fileInput.files[0];
  288.             var formData = new FormData();
  289.             formData.append('video', file);
  290.  
  291.             fetch('/predict_video', {
  292.                 method: 'POST',
  293.                 body: formData
  294.             })
  295.             .then(response => response.json())
  296.             .then(data => {
  297.                 var videoElement = document.getElementById('processedVideo');
  298.                 // 修改路径为正确的 Flask url_for 生成的路径
  299.                 videoElement.src = '{{ url_for("static", filename="") }}' + data.output_video_path;
  300.                 videoElement.style.display = 'block';
  301.                 var downloadButton = document.getElementById('downloadButton');
  302.                 downloadButton.style.display = 'block';
  303.             })
  304.             .catch(error => console.error('Error:', error));
  305.         }
  306.  
  307.         // 下载处理后的视频
  308.         function downloadProcessedVideo() {
  309.             var videoElement = document.getElementById('processedVideo');
  310.             var url = videoElement.src;
  311.             var a = document.createElement('a');
  312.             a.href = url;
  313.             a.download = 'processed_video.mp4';
  314.             document.body.appendChild(a);
  315.             a.click();
  316.             document.body.removeChild(a);
  317.         }
  318.  
  319.         // 关闭模态框
  320.         function closeModal() {
  321.             var modal = document.getElementById('myModal');
  322.             modal.style.display = 'none';
  323.         }
  324.     </script>
  325. </body>
  326. </html>

使用说明：

将index.html放入templates文件夹中

运行app.py

model = torch.hub.load("E:\\pythonProject2\\flaskProject\\yolov5-master", "custom", path='weight/'+new_filename, source="local")

注意此处加载模型路径更改为自己的

如果模型读取不到，csdn上有相关解决方法

此处的best.py上传路径也要更改成自己的

save_path = os.path.join("E:\\pythonProject2\\flaskProject\\weight\\", new_filename)

使用说明：

先上传本地的模型，上传成功后等待模型加载，再上传照片/视频

如有问题，可联系作者，随时讨论。

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号