📚 课程目录第4章 YOLOv5推理与目标检测可视化

4. YOLOv5推理与目标检测可视化

任务目标

加载训练完成的YOLOv5权重，对单张图片和视频流执行前向推理，在输出帧上绘制类别标签与置信度检测框，并保存可视化结果。

环境依赖

Python ≥ 3.8，安装：torch、opencv-python、numpy。若使用官方YOLOv5仓库，还需git clone https://github.com/ultralytics/yolov5并安装requirements.txt。本案例直接使用官方detect.py的封装接口，不修改模型结构。

关键完整代码

import cv2
import torch
import numpy as np

# ---------- 1. 加载模型 ----------
# 方式A：使用官方仓库（推荐，自动处理类别与预处理）
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.augmentations import letterbox

weights = 'yolov5s.pt'  # 或训练好的 runs/train/exp/weights/best.pt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights, device=device)
model.eval()

# 类别名称（COCO 80类，若为自定义模型请替换）
with open('data/coco.yaml') as f:
    names = [line.split(': ')[1].strip() for line in f if 'names:' in line][0].split(', ')
# 更稳妥：直接读取模型自带names
names = model.module.names if hasattr(model, 'module') else model.names

# ---------- 2. 图片推理 ----------
def infer_image(img_path, conf_thres=0.25, iou_thres=0.45):
    img0 = cv2.imread(img_path)
    img = letterbox(img0, new_shape=640)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR→RGB, HWC→CHW
    img = np.ascontiguousarray(img)
    img_tensor = torch.from_numpy(img).to(device).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0)

    with torch.no_grad():
        pred = model(img_tensor)[0]
    pred = non_max_suppression(pred, conf_thres, iou_thres)

    det = pred[0]
    if det is not None and len(det):
        det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], img0.shape).round()
        for *xyxy, conf, cls in det:
            x1, y1, x2, y2 = map(int, xyxy)
            label = f'{names[int(cls)]} {conf:.2f}'
            cv2.rectangle(img0, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img0, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    return img0

# ---------- 3. 视频推理 ----------
def infer_video(video_path, output_path):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        # 对每一帧执行与图片相同的预处理和推理
        img = letterbox(frame, new_shape=640)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img_tensor = torch.from_numpy(img).to(device).float() / 255.0
        img_tensor = img_tensor.unsqueeze(0)

        with torch.no_grad():
            pred = model(img_tensor)[0]
        pred = non_max_suppression(pred, 0.25, 0.45)

        det = pred[0]
        if det is not None and len(det):
            det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], frame.shape).round()
            for *xyxy, conf, cls in det:
                x1, y1, x2, y2 = map(int, xyxy)
                label = f'{names[int(cls)]} {conf:.2f}'
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, label, (x1, y1 - 5),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        writer.write(frame)
    cap.release()
    writer.release()

# ---------- 4. 调用入口 ----------
if __name__ == '__main__':
    # 图片推理
    result_img = infer_image('bus.jpg')
    cv2.imwrite('output_image.jpg', result_img)

    # 视频推理
    infer_video('test.mp4', 'output_video.mp4')

运行方式

将上述代码保存为infer_demo.py，放置在YOLOv5仓库根目录下（确保models、utils包可导入）。
准备测试图片bus.jpg（官方示例图）或任意含目标的图片；视频文件test.mp4。
执行：python infer_demo.py。

验证方法

图片输出：打开output_image.jpg，应看到目标周围有绿色矩形框，框上方标注“类别名置信度”（如“person 0.88”）。
视频输出：播放output_video.mp4，检测框应逐帧跟随目标移动，无卡顿或花屏。
判定通过条件：检测框位置与目标轮廓基本吻合，类别标签正确，置信度数值在0~1之间；视频帧率与原视频一致。

常见问题

导入错误：代码必须在YOLOv5仓库根目录运行，否则models、utils包无法找到。
自定义模型：若使用自己训练的权重，names列表需与训练时的data.yaml一致，否则标签显示错误。
GPU显存不足：将device改为'cpu'，或将letterbox的new_shape调小至320。