yolov5使用flask部署至前端實現照片\視頻識別功能

更新時間：2024年09月28日 09:36:28 作者：代碼與CPython

初學者在使用YOLO和Flask構建應用時,往往需要實現上傳圖片和視頻的識別功能,本文介紹了如何在Flask框架中實現這一功能,包括文件上傳、圖片放大查看、視頻識別以及識別后的文件下載,文中通過代碼介紹的非常詳細,需要的朋友可以參考下

前言

初學yolo flask時，需要此功能，Csdn、Github、B站找到許多代碼，效果并不滿意。

近期，再度嘗試，實現簡單功能。

實現功能：

上傳圖片并識別,可以點擊圖片放大查看

上傳視頻并識別

識別后的文件下載功能

效果圖如下

文件結構如下：

project/
static/

空

templates/
index.html

app.py

相關代碼：

app.py

import cv2
import numpy as np
import torch
from flask import Flask, request, jsonify, render_template
import base64
import os
from datetime import datetime

app = Flask(__name__)

# 全局變量：模型
model = None

# 提前加載模型
def load_model():
    global model
    model = torch.hub.load('', 'custom', path='yolov5s.pt', source='local')

# 路由處理圖片檢測請求
@app.route('/predict_image', methods=['POST'])
def predict_image():
    global model

    # 獲取圖像文件
    file = request.files['image']
    # 讀取圖像數據并轉換為RGB格式
    image_data = file.read()
    nparr = np.frombuffer(image_data, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)

    results = model(image)

    image = results.render()[0]

    # 將圖像轉換為 base64 編碼的字符串
    _, buffer = cv2.imencode('.png', image)
    image_str = base64.b64encode(buffer).decode('utf-8')

    # 獲取當前時間，并將其格式化為字符串
    current_time = datetime.now().strftime('%Y%m%d%H%M%S')
    # 構建保存路徑
    save_dir = 'static'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    filename, extension = os.path.splitext(file.filename)  # 獲取上傳文件的文件名和擴展名
    save_filename = f'{filename}_{current_time}{extension}'
    save_path = os.path.join(save_dir, save_filename)
    cv2.imwrite(save_path, image)

    return jsonify({'image': image_str})

# 函數用于在視頻幀上繪制檢測結果
def detect_objects(frame, model):
    results = model(frame)
    detections = results.xyxy[0].cpu().numpy()  # 獲取檢測結果

    # 在幀上繪制檢測結果
    for det in detections:
        # 獲取邊界框信息
        x1, y1, x2, y2, conf, class_id = det[:6]

        # 在幀上繪制邊界框
        cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)

        # 在幀上繪制類別和置信度
        label = f'{model.names[int(class_id)]} {conf:.2f}'
        cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    return frame

# 路由處理視頻檢測請求

@app.route("/predict_video", methods=["POST"])
def predict_video():
    global model

    # 從請求中獲取視頻文件
    video_file = request.files["video"]
    # 保存視頻到臨時文件
    temp_video_path = "temp_video.mp4"
    video_file.save(temp_video_path)

    # 逐幀讀取視頻
    video = cv2.VideoCapture(temp_video_path)

    # 獲取視頻的幀率和尺寸
    fps = video.get(cv2.CAP_PROP_FPS)
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 視頻寫入對象
    output_video_filename = f"output_video_{datetime.now().strftime('%Y%m%d%H%M%S')}.mp4"
    output_video_path = os.path.join("static", output_video_filename)
    fourcc = cv2.VideoWriter_fourcc(*"avc1")
    out_video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    # 逐幀處理視頻并進行目標檢測
    while True:
        ret, frame = video.read()
        if not ret:
            break

        # 進行目標檢測
        detection_result = detect_objects(frame, model)

        # 將處理后的幀寫入輸出視頻
        out_video.write(detection_result)

    # 釋放視頻對象
    video.release()
    out_video.release()

    return jsonify({"output_video_path": output_video_filename})

@app.route('/')
def index():
    return render_template('index.html')

# 初始加載模型
load_model()

if __name__ == '__main__':
    app.run(debug=True)

index.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Object Detection</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            margin: 0;
            padding: 0;
            background-color: #f3f3f3;
            display: flex;
            justify-content: center;
            align-items: center;
            height: 100vh;
            flex-direction: column;
        }

        #content {
            text-align: center;
            max-width: 820px;
            margin-top: 20px;
        }

        h1 {
            color: #333;
        }

        h2 {
            color: #666;
        }

        input[type="file"] {
            margin-bottom: 10px;
        }

        .media-container {
            display: flex;
            max-width: 100%;
            margin-bottom: 20px;
        }

        .media-container:first-child {
            margin-right: 20px; /* 在第一個容器的右側添加間隔 */
        }

        .media-container img,
        .media-container video {
            max-width: 100%;
            height: auto;
        }

        .original {
            width: 400px;
            overflow: hidden;
        }

        .processed {
            flex: 2; /* 右邊容器占據剩余空間 */
        }

        button {
            padding: 10px 20px;
            background-color: #007bff;
            color: #fff;
            border: none;
            border-radius: 5px;
            cursor: pointer;
            margin-bottom: 10px;
        }

        /* 新增樣式：模態(tài)框 */
        .modal {
            display: none; /* 默認隱藏 */
            position: fixed;
            z-index: 1;
            left: 0;
            top: 0;
            width: 100%;
            height: 100%;
            overflow: auto;
            background-color: rgba(0, 0, 0, 0.9); /* 半透明黑色背景 */
        }

        .modal-content {
            margin: auto;
            display: block;
            width: 80%;
            max-width: 800px;
            position: absolute;
            left: 50%;
            top: 50%;
            transform: translate(-50%, -50%);
            text-align: center; /* 居中顯示圖片 */
        }

        .close {
            color: #ccc;
            font-size: 36px;
            font-weight: bold;
            cursor: pointer;
            position: absolute;
            top: 10px;
            right: 10px;
        }

        .close:hover,
        .close:focus {
            color: #fff;
            text-decoration: none;
        }
        #downloadButton {
           padding: 10px 20px;
            background-color: #007bff;
            color: #fff;
            border: none;
            border-radius: 5px;
            cursor: pointer;
            margin-bottom: 10px;

        }

        /* 新增樣式：響應式圖片 */
        .modal-content img,
        .modal-content video {
            max-width: 100%;
            height: auto;
        }

    </style>
</head>
<body>

    <!-- 新增模態(tài)框 -->
    <div id="myModal" class="modal" onclick="closeModal()">
        <div class="modal-content" id="modalContent" onclick="stopPropagation(event)">
            <!-- 放大后的圖片或視頻將在這里顯示 -->
            <span class="close" onclick="closeModal()">&times;</span>
        </div>
    </div>

    <div id="content">
        <h1>照片/視頻檢測</h1>

        <!-- 上傳圖片 -->
        <h2>上傳圖片</h2>
        <input type="file" id="imageFile" accept="image/*" onchange="displaySelectedImage()">
        <button onclick="uploadImage()">上傳</button>
        <button id="downloadImageButton"  onclick="downloadProcessedImage()">下載</button>
        <br>
        <div class="media-container">
            <div class="original media-container" onclick="enlargeImage()">
                <img id="uploadedImage" src="#" alt="Uploaded Image" style="display:none;">
                <button id="zoomInButton" style="display:none;">Zoom In</button>
            </div>
            <div class="processed media-container" onclick="enlargeImage2()">
                <img id="processedImage" src="#" alt="Processed Image" style="display:none;">

            </div>
        </div>
        <br>

        <!-- 上傳視頻 -->
        <h2>上傳視頻</h2>
        <input type="file" id="videoFile" accept="video/mp4,video/x-m4v,video/*" onchange="displaySelectedVideo()">
        <button onclick="uploadVideo()">上傳</button>
        <button id="downloadButton" onclick="downloadProcessedVideo()">下載</button>
        <br>
        <div class="media-container">
            <div class="original media-container" >
                <video id="uploadedVideo" src="#" controls style="display:none;"></video>
            </div>
            <div class="processed media-container">
                <video id="processedVideo" controls style="display:none;"></video>

            </div>
        </div>
        <br>

    </div>

    <script>
         // 顯示選擇的權重文件

        // 顯示選擇的圖片并添加點擊放大功能
        function displaySelectedImage() {
            var fileInput = document.getElementById('imageFile');
            var file = fileInput.files[0];
            var imageElement = document.getElementById('uploadedImage');
            imageElement.src = URL.createObjectURL(file);
            imageElement.style.display = 'inline';
            document.getElementById('zoomInButton').style.display = 'inline';
        }

        // 顯示模態(tài)框并放大圖片
        function enlargeImage() {
            var modal = document.getElementById('myModal');
            var modalImg = document.getElementById('modalContent');
            var img = document.getElementById('uploadedImage');
            modal.style.display = 'block';
            modalImg.innerHTML = '<img src="' + img.src + '">';
        }
        // 顯示模態(tài)框并放大圖片
        function enlargeImage2() {
            var modal = document.getElementById('myModal');
            var modalImg = document.getElementById('modalContent');
            var img = document.getElementById('processedImage');
            modal.style.display = 'block';
            modalImg.innerHTML = '<img src="' + img.src + '">';
        }


        // 顯示選擇的視頻并添加點擊放大功能
        function displaySelectedVideo() {
            var fileInput = document.getElementById('videoFile');
            var file = fileInput.files[0];
            var videoElement = document.getElementById('uploadedVideo');
            videoElement.src = URL.createObjectURL(file);
            videoElement.style.display = 'block';
        }


        // 上傳圖片并向后端發(fā)送請求
        function uploadImage() {
            var fileInput = document.getElementById('imageFile');
            var file = fileInput.files[0];
            var formData = new FormData();
            formData.append('image', file);

            fetch('/predict_image', {
                method: 'POST',
                body: formData
            })
            .then(response => response.json())
            .then(data => {
                var imageElement = document.getElementById('processedImage');
                imageElement.src = 'data:image/png;base64,' + data.image;
                imageElement.style.display = 'inline';
                document.getElementById('downloadImageButton').style.display = 'inline';
            })
            .catch(error => console.error('Error:', error));
        }

        // 下載處理后的圖片
        function downloadProcessedImage() {
            var imageElement = document.getElementById('processedImage');
            var url = imageElement.src;
            var a = document.createElement('a');
            a.href = url;
            a.download = 'processed_image.png';
            document.body.appendChild(a);
            a.click();
            document.body.removeChild(a);
        }

        // 上傳視頻并向后端發(fā)送請求
        function uploadVideo() {
            var fileInput = document.getElementById('videoFile');
            var file = fileInput.files[0];
            var formData = new FormData();
            formData.append('video', file);

            fetch('/predict_video', {
                method: 'POST',
                body: formData
            })
            .then(response => response.json())
            .then(data => {
                var videoElement = document.getElementById('processedVideo');
                // 修改路徑為正確的 Flask url_for 生成的路徑
                videoElement.src = '{{ url_for("static", filename="") }}' + data.output_video_path;
                videoElement.style.display = 'block';
                var downloadButton = document.getElementById('downloadButton');
                downloadButton.style.display = 'block';
            })
            .catch(error => console.error('Error:', error));
        }

        // 下載處理后的視頻
        function downloadProcessedVideo() {
            var videoElement = document.getElementById('processedVideo');
            var url = videoElement.src;
            var a = document.createElement('a');
            a.href = url;
            a.download = 'processed_video.mp4';
            document.body.appendChild(a);
            a.click();
            document.body.removeChild(a);
        }

        // 關閉模態(tài)框
        function closeModal() {
            var modal = document.getElementById('myModal');
            modal.style.display = 'none';
        }
    </script>
</body>
</html>

使用說明：

index.html放入templates文件夾中

運行app.py

注：此處加載模型路徑更改為自己的

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文件復制到y(tǒng)olov5-master文件夾中，修改model路徑

model = torch.hub.load('yolov5-master', 'custom', path='yolov5s.pt', source='local')

總結

到此這篇關于yolov5使用flask部署至前端實現照片\視頻識別功能的文章就介紹到這了,更多相關yolov5 flask部署前端照片視頻識別內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！

您可能感興趣的文章:

YOLOv5部署到web端詳細過程(flask+js簡單易懂)

Python最長回文子串問題
這篇文章主要介紹了Python最長回文子串問題，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教
2022-11-11
python局域網ip掃描示例分享
這篇文章主要介紹了python局域網ip掃描示例,需要的朋友可以參考下
2014-04-04
Python使用Beautiful?Soup(BS4)庫解析HTML和XML
這篇文章介紹了Python使用Beautiful?Soup(BS4)庫解析HTML和XML的方法，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧
2022-06-06
10個Python自動化辦公的腳本分享
在日常辦公中,我們常常會被繁瑣、重復的任務占據大量時間,本文為大家分享了10個實用的?Python?自動化辦公案例及源碼,希望對大家有所幫助
2025-02-02
Python連接SQLServer2000的方法詳解
這篇文章主要介紹了Python連接SQLServer2000的方法,結合實例形式分析了Python實現數據庫連接過程中所遇到的常見問題與相關注意事項,需要的朋友可以參考下
2017-04-04
最小公倍數Python實現的方法例子
這篇文章介紹了兩種計算最小公倍數的方法：輾轉相除法（歐幾里德法）和相減法（更相減損法）,這兩種方法分別基于求最大公約數的不同原理,需要的朋友可以參考下
2024-11-11
matplotlib繪制折線圖的基本配置(萬能模板案例)
折線圖可以很方便的看出數據的對比，本文主要介紹了matplotlib繪制折線圖的基本配置(萬能模板案例)，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧
2022-04-04
Python 迭代器與生成器實例詳解
這篇文章主要介紹了Python 迭代器與生成器實例詳解的相關資料,需要的朋友可以參考下
2017-05-05
python字典序問題實例
這篇文章主要介紹了python字典序問題,是字符串操作一個比較典型的應用,需要的朋友可以參考下
2014-09-09
NumPy中掩碼數組的操作
本文主要介紹了NumPy中掩碼數組的操作，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧
2023-03-03