yolov5 6.0中对摄像头视频帧进行推理的代码

以下是使用 YOLOv5 6.0 对摄像头视频帧进行推理的示例代码：

import cv2
import torch
import numpy as np
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords, plot_one_box
from utils.torch_utils import select_device

# 加载模型
weights = 'yolov5s.pt'  # 模型权重文件路径
device = select_device('')  # 自动选择可用的设备
model = attempt_load(weights, map_location=device)
stride = int(model.stride.max())  # 计算特征图的步长

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    ret, img = cap.read()  # 读取摄像头视频帧
    if not ret:
        break

    # 图像预处理
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR 转 RGB
    img = cv2.resize(img, (640, 640), interpolation=cv2.INTER_LINEAR)  # 缩放到模型输入大小
    img = img.transpose(2, 0, 1)  # HWC 转 CHW
    img = np.ascontiguousarray(img)  # 转换为连续内存

    # 模型推理
    img = torch.from_numpy(img).to(device).float() / 255.0
    img = img.unsqueeze(0)
    pred = model(img)[0]
    pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)

    # 后处理
    for det in pred:
        if det is not None and len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[2:]).round()
            for *xyxy, conf, cls in reversed(det):
                label = f'{model.names[int(cls)]} {conf:.2f}'
                plot_one_box(xyxy, img, label=label, color=(0, 255, 0), line_thickness=2)

    # 显示图像
    img = img.squeeze(0).permute(1, 2, 0).cpu().numpy()  # 转回 HWC
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # RGB 转 BGR
    cv2.imshow('YOLOv5', img)

    if cv2.waitKey(1) == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

代码中首先加载 YOLOv5 模型，然后打开摄像头并循环读取视频帧。接着，对每个视频帧进行预处理，包括 BGR 转 RGB、缩放到模型输入大小、HWC 转 CHW 和转换为连续内存等操作。然后，将视频帧输入模型进行推理，得到预测结果，并进行后处理，包括特征图上的非极大值抑制、坐标缩放、画框和写标签等操作。最后，将处理后的图像显示出来，并在按下 q 键时退出循环并释放资源。