如何将Flask+YOLOv8服务从图片检测改为支持.mov视频检测？

后端代码修改（object_detector.py） #

首先安装OpenCV依赖：

pip install opencv-python

修改后的完整代码：

from ultralytics import YOLO
from flask import request, Flask, jsonify, send_file
from waitress import serve
from PIL import Image
import cv2
import tempfile
import os

app = Flask(__name__)

# 提前加载模型，避免重复加载损耗性能
model = YOLO("bestnano.pt")

@app.route("/")
def root():
    with open("index.html") as file:
        return file.read()

@app.route("/detect", methods=["POST"])
def detect():
    buf = request.files["media_file"]
    filename = buf.filename
    # 根据后缀判断文件类型
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        boxes = detect_objects_on_image(buf.stream)
        return jsonify({"type": "image", "data": boxes})
    elif filename.lower().endswith(('.mov', '.mp4', '.avi')):
        # 生成临时输出视频文件
        temp_output = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
        temp_output.close()
        process_video(buf.stream, temp_output.name)
        return send_file(temp_output.name, mimetype='video/mp4', as_attachment=False, download_name='detected_video.mp4')
    else:
        return jsonify({"error": "不支持的文件格式"}), 400

def detect_objects_on_image(buf):
    results = model.predict(Image.open(buf))
    result = results[0]
    output = []
    for box in result.boxes:
        x1, y1, x2, y2 = [round(x) for x in box.xyxy[0].tolist()]
        class_id = box.cls[0].item()
        prob = round(box.conf[0].item(), 2)
        output.append([x1, y1, x2, y2, result.names[class_id], prob])
    return output

def process_video(video_stream, output_path):
    # 保存上传的视频到临时文件
    temp_input = tempfile.NamedTemporaryFile(suffix='.mov', delete=False)
    temp_input.write(video_stream.read())
    temp_input.close()

    # 初始化视频读取器
    cap = cv2.VideoCapture(temp_input.name)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 初始化视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 转换帧格式用于YOLO检测
        pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        results = model.predict(pil_frame)
        result = results[0]

        # 在帧上绘制检测框和标签
        for box in result.boxes:
            x1, y1, x2, y2 = [round(x) for x in box.xyxy[0].tolist()]
            class_id = box.cls[0].item()
            label = f"{result.names[class_id]} {round(box.conf[0].item(), 2)}"
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 3)
            cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        out.write(frame)

    # 释放资源并清理临时文件
    cap.release()
    out.release()
    os.unlink(temp_input.name)

serve(app, host='0.0.0.0', port=8080)

核心修改点：

• 提前加载YOLO模型，提升请求处理效率
• 扩展/detect端点，支持图片和视频两种文件类型的分支处理
• 新增视频处理逻辑：保存临时输入视频、逐帧检测并绘制标注、生成带检测框的输出视频

前端代码修改（index.html） #

<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <title>YOLOv8 交通信号灯检测</title>
    <style>
        canvas, video {
            display:block;
            border: 1px solid black;
            margin-top:10px;
            max-width: 800px;
        }
    </style>
</head>
<body>
    <input id="uploadInput" type="file" accept=".png,.jpg,.jpeg,.mov,.mp4,.avi"/>
    <canvas id="imageCanvas"></canvas>
    <video id="videoPlayer" controls></video>

    <script>
        const input = document.getElementById("uploadInput");
        const canvas = document.getElementById("imageCanvas");
        const videoPlayer = document.getElementById("videoPlayer");

        input.addEventListener("change",async(event) => {
            // 隐藏上一次的展示内容
            canvas.style.display = "none";
            videoPlayer.style.display = "none";

            const file = event.target.files[0];
            const data = new FormData();
            data.append("media_file", file);
            const response = await fetch("/detect",{
                method:"post",
                body:data
            });

            if (response.headers.get('Content-Type').includes('application/json')) {
                // 处理图片检测结果
                const result = await response.json();
                if (result.type === "image") {
                    canvas.style.display = "block";
                    draw_image_and_boxes(file, result.data);
                } else {
                    alert(result.error);
                }
            } else {
                // 处理视频检测结果
                videoPlayer.style.display = "block";
                const blob = await response.blob();
                videoPlayer.src = URL.createObjectURL(blob);
                videoPlayer.play();
            }
        })

        function draw_image_and_boxes(file,boxes) {
            const img = new Image()
            img.src = URL.createObjectURL(file);
            img.onload = () => {
                canvas.width = img.width;
                canvas.height = img.height;
                const ctx = canvas.getContext("2d");
                ctx.drawImage(img,0,0);
                ctx.strokeStyle = "#00FF00";
                ctx.lineWidth = 3;
                ctx.font = "18px serif";
                boxes.forEach(([x1,y1,x2,y2,label,prob]) => {
                    const fullLabel = `${label} ${prob}`;
                    ctx.strokeRect(x1,y1,x2-x1,y2-y1);
                    ctx.fillStyle = "#00ff00";
                    const width = ctx.measureText(fullLabel).width;
                    ctx.fillRect(x1,y1,width+10,25);
                    ctx.fillStyle = "#000000";
                    ctx.fillText(fullLabel, x1, y1+18);
                });
            }
        }
    </script>
</body>
</html>

核心修改点：

• 扩展文件选择器的支持格式，添加MOV等视频类型
• 新增video标签用于播放处理后的视频
• 区分图片和视频的处理逻辑：图片用canvas绘制标注，视频直接播放后端返回的带检测框的视频
• 优化图片标注显示，增加检测概率信息

内容的提问来源于stack exchange，提问作者hejjhej