2. YOLO视频目标检测基础

任务目标

使用OpenCV的dnn模块加载YOLOv4模型，对视频文件逐帧执行目标检测，在窗口中实时显示带边界框和类别标签的检测结果，同时将处理后的帧写入输出视频文件。

环境依赖

• OpenCV 4.x（含dnn模块）
• YOLOv4权重文件 yolov4.weights、配置文件 yolov4.cfg、类别文件 coco.names（可从Darknet官方仓库获取）
• 输入视频文件（如 input.mp4）

完整实现代码

import cv2
import numpy as np

# ========== 配置 ==========
weights_path = "yolov4.weights"
config_path = "yolov4.cfg"
names_path = "coco.names"
input_video = "input.mp4"
output_video = "output.avi"
conf_threshold = 0.5
nms_threshold = 0.4
input_size = (416, 416)

# ========== 加载模型与类别 ==========
net = cv2.dnn.readNet(weights_path, config_path)
with open(names_path, "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]

# ========== 视频读写 ==========
cap = cv2.VideoCapture(input_video)
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(*"XVID")
writer = cv2.VideoWriter(output_video, fourcc, fps, (width, height))

# ========== 检测函数 ==========
def detect_objects(frame):
    h, w = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(frame, 1/255.0, input_size, swapRB=True, crop=False)
    net.setInput(blob)
    outputs = net.forward(output_layers)

    boxes, confidences, class_ids = [], [], []
    for output in outputs:
        for detection in output:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > conf_threshold:
                cx, cy, bw, bh = detection[:4] * np.array([w, h, w, h])
                x = int(cx - bw / 2)
                y = int(cy - bh / 2)
                boxes.append([x, y, int(bw), int(bh)])
                confidences.append(float(confidence))
                class_ids.append(class_id)

    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
    return boxes, confidences, class_ids, indices

# ========== 主循环 ==========
while True:
    ret, frame = cap.read()
    if not ret:
        break

    boxes, confidences, class_ids, indices = detect_objects(frame)

    if len(indices) > 0:
        for i in indices.flatten():
            x, y, bw, bh = boxes[i]
            label = f"{classes[class_ids[i]]}: {confidences[i]:.2f}"
            cv2.rectangle(frame, (x, y), (x + bw, y + bh), (0, 255, 0), 2)
            cv2.putText(frame, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    cv2.imshow("YOLO Video Detection", frame)
    writer.write(frame)

    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

cap.release()
writer.release()
cv2.destroyAllWindows()

运行方式

1. 将代码保存为 yolo_video.py，与模型文件、类别文件及输入视频放在同一目录。
2. 终端执行：python yolo_video.py

验证方法

• 程序启动后弹出窗口，视频帧中检测到的目标被绿色边界框框出，框上方显示类别名称和置信度。
• 按 q 键可提前终止；视频播放完毕自动结束。
• 检查当前目录下生成的 output.avi，用播放器打开，确认其内容与窗口显示一致，且时长与原视频相同。