如何用自定义训练的YOLOv7逐帧处理视频并检测目标？

核心原因与解决方案 #

直接传入OpenCV帧结果异常，本质是OpenCV读取的BGR格式帧与YOLOv7要求的RGB输入格式不匹配，且缺少模型需要的标准化预处理步骤。无需转换为TensorFlow/ONNX格式，用PyTorch就能直接实现需求，只需对齐模型输入格式即可。

完整实现流程 #

1. 加载自定义训练的YOLOv7模型 #

先加载你的权重文件，切换到推理模式：

import torch
from models.experimental import attempt_load

# 替换为你的自定义权重路径
weights = "./runs/train/exp/weights/best.pt"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载模型并设置为推理模式
model = attempt_load(weights, map_location=device)
model.eval()

2. 定义帧预处理函数 #

将OpenCV的BGR帧转换成模型可接受的张量，保证和训练时的预处理逻辑一致：

import cv2
import numpy as np
from utils.general import letterbox

def preprocess(frame, img_size=640):
    # 1. OpenCV读的是BGR，转成YOLOv7需要的RGB
    img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 2. 按比例调整尺寸，补黑边（和YOLOv7自带预处理对齐）
    img_resized, _, _ = letterbox(img_rgb, new_shape=img_size)
    # 3. 转成PyTorch张量，调整维度顺序，归一化
    img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).float().to(device)
    img_tensor /= 255.0
    # 4. 增加batch维度
    img_tensor = img_tensor.unsqueeze(0)
    # 返回预处理后的张量和原始帧尺寸（用于后续坐标缩放）
    return img_tensor, frame.shape[:2]

3. 逐帧处理视频并解析结果 #

用OpenCV读取视频，逐帧预处理后传入模型，解析得到目标坐标、置信度和类别：

from utils.general import non_max_suppression, scale_coords

# 替换为你的视频路径
video_path = "./test_video.mp4"
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 预处理当前帧
    img_tensor, original_shape = preprocess(frame)
    
    # 模型推理，关闭梯度计算提升速度
    with torch.no_grad():
        pred = model(img_tensor)[0]
    
    # 非极大值抑制，过滤重复检测框
    pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
    
    # 解析每一个检测结果
    for det in pred:
        if len(det):
            # 将检测框从调整后的尺寸缩放回原始帧尺寸
            det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], original_shape).round()
            
            # 遍历每个检测目标，提取数据
            for *xyxy, conf, cls in det:
                x1, y1, x2, y2 = map(int, xyxy)
                confidence = float(conf)
                class_id = int(cls)
                
                # 这里写你自定义的处理逻辑，比如存储结果、分析数据等
                print(f"类别ID: {class_id}, 置信度: {confidence:.2f}, 坐标: ({x1},{y1})-({x2},{y2})")
                
                # 可选：在帧上绘制检测框（测试用）
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, f"{class_id} {confidence:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
    
    # 可选：显示处理后的帧
    cv2.imshow("Detection", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

关键注意事项 #

• 格式转换：必须将OpenCV的BGR转为RGB，这是之前结果异常的核心原因
• 预处理对齐：使用YOLOv7自带的letterbox函数调整尺寸，保证和训练时的输入一致
• 推理模式：模型必须设置为eval()，关闭训练时的批量归一化、dropout逻辑，避免结果不稳定
• 坐标缩放：检测框基于调整后的图像尺寸，必须用scale_coords缩放回原始帧尺寸，才能得到正确位置

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