yolov5使用flask部署至前端,实现照片\视频识别

打印 上一主题 下一主题

主题 1555|帖子 1555|积分 4665

马上注册,结交更多好友,享用更多功能,让你轻松玩转社区。

您需要 登录 才可以下载或查看,没有账号?立即注册

x
初学yolo flask时,必要此功能,Csdn、Github、B站找到许多代码,效果并不满意。

近期,再度尝试,实现简单功能。



实现功能:


  • 上传图片并识别,可以点击图片放大查看
  • 上传视频并识别
  • 识别后的文件下载功能

效果图如上

文件结构如下:
   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. app = Flask(__name__)
  9. # 全局变量:模型
  10. model = None
  11. # 提前加载模型
  12. def load_model():
  13.     global model
  14.     model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')
  15. # 路由处理图片检测请求
  16. @app.route('/predict_image', methods=['POST'])
  17. def predict_image():
  18.     global model
  19.     # 获取图像文件
  20.     file = request.files['image']
  21.     # 读取图像数据并转换为RGB格式
  22.     image_data = file.read()
  23.     nparr = np.frombuffer(image_data, np.uint8)
  24.     image = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
  25.     results = model(image)
  26.     image = results.render()[0]
  27.     # 将图像转换为 base64 编码的字符串
  28.     _, buffer = cv2.imencode('.png', image)
  29.     image_str = base64.b64encode(buffer).decode('utf-8')
  30.     # 获取当前时间,并将其格式化为字符串
  31.     current_time = datetime.now().strftime('%Y%m%d%H%M%S')
  32.     # 构建保存路径
  33.     save_dir = 'static'
  34.     if not os.path.exists(save_dir):
  35.         os.makedirs(save_dir)
  36.     filename, extension = os.path.splitext(file.filename)  # 获取上传文件的文件名和扩展名
  37.     save_filename = f'{filename}_{current_time}{extension}'
  38.     save_path = os.path.join(save_dir, save_filename)
  39.     cv2.imwrite(save_path, image)
  40.     return jsonify({'image': image_str})
  41. # 函数用于在视频帧上绘制检测结果
  42. def detect_objects(frame, model):
  43.     results = model(frame)
  44.     detections = results.xyxy[0].cpu().numpy()  # 获取检测结果
  45.     # 在帧上绘制检测结果
  46.     for det in detections:
  47.         # 获取边界框信息
  48.         x1, y1, x2, y2, conf, class_id = det[:6]
  49.         # 在帧上绘制边界框
  50.         cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
  51.         # 在帧上绘制类别和置信度
  52.         label = f'{model.names[int(class_id)]} {conf:.2f}'
  53.         cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
  54.     return frame
  55. # 路由处理视频检测请求
  56. @app.route("/predict_video", methods=["POST"])
  57. def predict_video():
  58.     global model
  59.     # 从请求中获取视频文件
  60.     video_file = request.files["video"]
  61.     # 保存视频到临时文件
  62.     temp_video_path = "temp_video.mp4"
  63.     video_file.save(temp_video_path)
  64.     # 逐帧读取视频
  65.     video = cv2.VideoCapture(temp_video_path)
  66.     # 获取视频的帧率和尺寸
  67.     fps = video.get(cv2.CAP_PROP_FPS)
  68.     width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
  69.     height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
  70.     # 视频写入对象
  71.     output_video_filename = f"output_video_{datetime.now().strftime('%Y%m%d%H%M%S')}.mp4"
  72.     output_video_path = os.path.join("static", output_video_filename)
  73.     fourcc = cv2.VideoWriter_fourcc(*"avc1")
  74.     out_video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
  75.     # 逐帧处理视频并进行目标检测
  76.     while True:
  77.         ret, frame = video.read()
  78.         if not ret:
  79.             break
  80.         # 进行目标检测
  81.         detection_result = detect_objects(frame, model)
  82.         # 将处理后的帧写入输出视频
  83.         out_video.write(detection_result)
  84.     # 释放视频对象
  85.     video.release()
  86.     out_video.release()
  87.     return jsonify({"output_video_path": output_video_filename})
  88. @app.route('/')
  89. def index():
  90.     return render_template('index.html')
  91. # 初始加载模型
  92. load_model()
  93. if __name__ == '__main__':
  94.     app.run(debug=True)
复制代码
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.         #content {
  20.             text-align: center;
  21.             max-width: 820px;
  22.             margin-top: 20px;
  23.         }
  24.         h1 {
  25.             color: #333;
  26.         }
  27.         h2 {
  28.             color: #666;
  29.         }
  30.         input[type="file"] {
  31.             margin-bottom: 10px;
  32.         }
  33.         .media-container {
  34.             display: flex;
  35.             max-width: 100%;
  36.             margin-bottom: 20px;
  37.         }
  38.         .media-container:first-child {
  39.             margin-right: 20px; /* 在第一个容器的右侧添加间隔 */
  40.         }
  41.         .media-container img,
  42.         .media-container video {
  43.             max-width: 100%;
  44.             height: auto;
  45.         }
  46.         .original {
  47.             width: 400px;
  48.             overflow: hidden;
  49.         }
  50.         .processed {
  51.             flex: 2; /* 右边容器占据剩余空间 */
  52.         }
  53.         button {
  54.             padding: 10px 20px;
  55.             background-color: #007bff;
  56.             color: #fff;
  57.             border: none;
  58.             border-radius: 5px;
  59.             cursor: pointer;
  60.             margin-bottom: 10px;
  61.         }
  62.         /* 新增样式:模态框 */
  63.         .modal {
  64.             display: none; /* 默认隐藏 */
  65.             position: fixed;
  66.             z-index: 1;
  67.             left: 0;
  68.             top: 0;
  69.             width: 100%;
  70.             height: 100%;
  71.             overflow: auto;
  72.             background-color: rgba(0, 0, 0, 0.9); /* 半透明黑色背景 */
  73.         }
  74.         .modal-content {
  75.             margin: auto;
  76.             display: block;
  77.             width: 80%;
  78.             max-width: 800px;
  79.             position: absolute;
  80.             left: 50%;
  81.             top: 50%;
  82.             transform: translate(-50%, -50%);
  83.             text-align: center; /* 居中显示图片 */
  84.         }
  85.         .close {
  86.             color: #ccc;
  87.             font-size: 36px;
  88.             font-weight: bold;
  89.             cursor: pointer;
  90.             position: absolute;
  91.             top: 10px;
  92.             right: 10px;
  93.         }
  94.         .close:hover,
  95.         .close:focus {
  96.             color: #fff;
  97.             text-decoration: none;
  98.         }
  99.         #downloadButton {
  100.            padding: 10px 20px;
  101.             background-color: #007bff;
  102.             color: #fff;
  103.             border: none;
  104.             border-radius: 5px;
  105.             cursor: pointer;
  106.             margin-bottom: 10px;
  107.         }
  108.         /* 新增样式:响应式图片 */
  109.         .modal-content img,
  110.         .modal-content video {
  111.             max-width: 100%;
  112.             height: auto;
  113.         }
  114.     </style>
  115. </head>
  116. <body>
  117.     <!-- 新增模态框 -->
  118.     <div id="myModal" class="modal" onclick="closeModal()">
  119.         <div class="modal-content" id="modalContent" onclick="stopPropagation(event)">
  120.             <!-- 放大后的图片或视频将在这里显示 -->
  121.             <span class="close" onclick="closeModal()">&times;</span>
  122.         </div>
  123.     </div>
  124.     <div id="content">
  125.         <h1>照片/视频检测</h1>
  126.         <!-- 上传图片 -->
  127.         <h2>上传图片</h2>
  128.         <input type="file" id="imageFile" accept="image/*" onchange="displaySelectedImage()">
  129.         <button onclick="uploadImage()">上传</button>
  130.         <button id="downloadImageButton"  onclick="downloadProcessedImage()">下载</button>
  131.         <br>
  132.         <div class="media-container">
  133.             <div class="original media-container" onclick="enlargeImage()">
  134.                 <img id="uploadedImage" src="#" alt="Uploaded Image" style="display:none;">
  135.                 <button id="zoomInButton" style="display:none;">Zoom In</button>
  136.             </div>
  137.             <div class="processed media-container" onclick="enlargeImage2()">
  138.                 <img id="processedImage" src="#" alt="Processed Image" style="display:none;">
  139.             </div>
  140.         </div>
  141.         <br>
  142.         <!-- 上传视频 -->
  143.         <h2>上传视频</h2>
  144.         <input type="file" id="videoFile" accept="video/mp4,video/x-m4v,video/*" onchange="displaySelectedVideo()">
  145.         <button onclick="uploadVideo()">上传</button>
  146.         <button id="downloadButton" onclick="downloadProcessedVideo()">下载</button>
  147.         <br>
  148.         <div class="media-container">
  149.             <div class="original media-container" >
  150.                 <video id="uploadedVideo" src="#" controls style="display:none;"></video>
  151.             </div>
  152.             <div class="processed media-container">
  153.                 <video id="processedVideo" controls style="display:none;"></video>
  154.             </div>
  155.         </div>
  156.         <br>
  157.     </div>
  158.     <script>
  159.          // 显示选择的权重文件
  160.         // 显示选择的图片并添加点击放大功能
  161.         function displaySelectedImage() {
  162.             var fileInput = document.getElementById('imageFile');
  163.             var file = fileInput.files[0];
  164.             var imageElement = document.getElementById('uploadedImage');
  165.             imageElement.src = URL.createObjectURL(file);
  166.             imageElement.style.display = 'inline';
  167.             document.getElementById('zoomInButton').style.display = 'inline';
  168.         }
  169.         // 显示模态框并放大图片
  170.         function enlargeImage() {
  171.             var modal = document.getElementById('myModal');
  172.             var modalImg = document.getElementById('modalContent');
  173.             var img = document.getElementById('uploadedImage');
  174.             modal.style.display = 'block';
  175.             modalImg.innerHTML = '<img src="' + img.src + '">';
  176.         }
  177.         // 显示模态框并放大图片
  178.         function enlargeImage2() {
  179.             var modal = document.getElementById('myModal');
  180.             var modalImg = document.getElementById('modalContent');
  181.             var img = document.getElementById('processedImage');
  182.             modal.style.display = 'block';
  183.             modalImg.innerHTML = '<img src="' + img.src + '">';
  184.         }
  185.         // 显示选择的视频并添加点击放大功能
  186.         function displaySelectedVideo() {
  187.             var fileInput = document.getElementById('videoFile');
  188.             var file = fileInput.files[0];
  189.             var videoElement = document.getElementById('uploadedVideo');
  190.             videoElement.src = URL.createObjectURL(file);
  191.             videoElement.style.display = 'block';
  192.         }
  193.         // 上传图片并向后端发送请求
  194.         function uploadImage() {
  195.             var fileInput = document.getElementById('imageFile');
  196.             var file = fileInput.files[0];
  197.             var formData = new FormData();
  198.             formData.append('image', file);
  199.             fetch('/predict_image', {
  200.                 method: 'POST',
  201.                 body: formData
  202.             })
  203.             .then(response => response.json())
  204.             .then(data => {
  205.                 var imageElement = document.getElementById('processedImage');
  206.                 imageElement.src = 'data:image/png;base64,' + data.image;
  207.                 imageElement.style.display = 'inline';
  208.                 document.getElementById('downloadImageButton').style.display = 'inline';
  209.             })
  210.             .catch(error => console.error('Error:', error));
  211.         }
  212.         // 下载处理后的图片
  213.         function downloadProcessedImage() {
  214.             var imageElement = document.getElementById('processedImage');
  215.             var url = imageElement.src;
  216.             var a = document.createElement('a');
  217.             a.href = url;
  218.             a.download = 'processed_image.png';
  219.             document.body.appendChild(a);
  220.             a.click();
  221.             document.body.removeChild(a);
  222.         }
  223.         // 上传视频并向后端发送请求
  224.         function uploadVideo() {
  225.             var fileInput = document.getElementById('videoFile');
  226.             var file = fileInput.files[0];
  227.             var formData = new FormData();
  228.             formData.append('video', file);
  229.             fetch('/predict_video', {
  230.                 method: 'POST',
  231.                 body: formData
  232.             })
  233.             .then(response => response.json())
  234.             .then(data => {
  235.                 var videoElement = document.getElementById('processedVideo');
  236.                 // 修改路径为正确的 Flask url_for 生成的路径
  237.                 videoElement.src = '{{ url_for("static", filename="") }}' + data.output_video_path;
  238.                 videoElement.style.display = 'block';
  239.                 var downloadButton = document.getElementById('downloadButton');
  240.                 downloadButton.style.display = 'block';
  241.             })
  242.             .catch(error => console.error('Error:', error));
  243.         }
  244.         // 下载处理后的视频
  245.         function downloadProcessedVideo() {
  246.             var videoElement = document.getElementById('processedVideo');
  247.             var url = videoElement.src;
  248.             var a = document.createElement('a');
  249.             a.href = url;
  250.             a.download = 'processed_video.mp4';
  251.             document.body.appendChild(a);
  252.             a.click();
  253.             document.body.removeChild(a);
  254.         }
  255.         // 关闭模态框
  256.         function closeModal() {
  257.             var modal = document.getElementById('myModal');
  258.             modal.style.display = 'none';
  259.         }
  260.     </script>
  261. </body>
  262. </html>
复制代码
使用分析:
index.html放入templates文件夹中
运行app.py
注:此处加载模子路径更改为自己的
  
  1. model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')
复制代码
如果模子读取不到,表现

   FileNotFoundError: [Errno 2] No such file or directory: 'hubconf.py'
  去yolov5官网,下载yolov5-master到项目文件夹

并将yolov5s.pt文件复制到yolov5-master文件夹中,修改model路径
  
  1. model = torch.hub.load('yolov5-master', 'custom', path='yolov5s.pt', source='local')
复制代码
如有问题,可接洽作者。

免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!更多信息从访问主页:qidao123.com:ToB企服之家,中国第一个企服评测及商务社交产业平台。
回复

使用道具 举报

0 个回复

倒序浏览

快速回复

您需要登录后才可以回帖 登录 or 立即注册

本版积分规则

冬雨财经

论坛元老
这个人很懒什么都没写!
快速回复 返回顶部 返回列表