利用摄像头输入读取视频的yolov5.pt推理测试程序

“拷贝即用” 的 YOLOv5 摄像头推理脚本：

• 支持 USB 摄像头 / RTSP / 本地视频文件
• 按 q 退出，按 s 保存当前画面到 runs/detect_cam/
• 自带 FPS 计数，分辨率可调
• 依赖：已安装 torch + opencv-python + yolov5 即可运行

---

📁 文件： detect_cam.py

#!/usr/bin/env python3
"""
YOLOv5 摄像头/视频实时推理
usage:
    python detect_cam.py --weights yolov5s.pt --source 0        # USB摄像头
    python detect_cam.py --weights yolov5s.pt --source rtsp://xxx
    python detect_cam.py --weights yolov5s.pt --source video.mp4
"""
import argparse
import time
from pathlib import Path

import cv2
import torch
from numpy import ndarray

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 根目录
print(f"ROOT:{ROOT}")

# ====== 1. 载入 yolov5 ======
model = torch.hub.load(ROOT, 'custom', path="yolov5s.pt", source='local')  # 本地 repo
model.eval()

# ====== 2. 参数 ======
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='yolov5s.pt', help='模型路径')
parser.add_argument('--source', type=str, default='0', help='0/1/2... 或 rtsp:// 或 *.mp4')
parser.add_argument('--img-size', type=int, default=640, help='推理尺寸')
parser.add_argument('--conf-thres', type=float, default=0.25, help='置信度阈值')
parser.add_argument('--save-dir', type=str, default='runs/detect_cam', help='保存目录')
opt = parser.parse_args()

# 数字则转 int
source = int(opt.source) if opt.source.isdigit() else opt.source
cap = cv2.VideoCapture(source)
assert cap.isOpened(), f'打不开摄像头/视频: {source}'

# 保存目录
save_path = Path(opt.save_dir)
save_path.mkdir(parents=True, exist_ok=True)
print(f'按 q 退出，按 s 保存截图到 {save_path}')

# ====== 3. 颜色 + FPS ======
names = model.names
colors = [(36, 255, 12), (255, 0, 0), (0, 255, 255)] * len(names)

# ====== 4. 推理函数 ======
@torch.no_grad()
def detect(frame: ndarray):
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = model(img, size=opt.img_size)
    return results

# ====== 5. 主循环 ======
frame_cnt = 0
while True:
    ret, frame = cap.read()
    if not ret:
        print('读取失败，退出')
        break

    t0 = time.time()
    results = detect(frame)          # yolov5 返回 pandas 格式
    fps = 1 / (time.time() - t0 + 1e-6)

    # 画框
    for *xyxy, conf, cls in results.xyxy[0].cpu().numpy():
        if conf < opt.conf_thres:
            continue
        x1, y1, x2, y2 = map(int, xyxy)
        label = f'{names[int(cls)]} {conf:.2f}'
        color = colors[int(cls) % len(colors)]
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, color, 2)

    # 显示 FPS
    cv2.putText(frame, f'FPS:{fps:.1f}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                1, (0, 0, 255), 2)
    cv2.imshow('YOLOv5-Cam', frame)

    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):
        break
    elif key == ord('s'):  # 截图
        f = save_path / f'capture_{int(time.time())}.jpg'
        cv2.imwrite(str(f), frame)
        print(f'已保存 {f}')

cap.release()
cv2.destroyAllWindows()

---

🔧 依赖安装

pip install torch torchvision opencv-python

---

🚀 运行示例

# USB 摄像头（默认 0）
python detect_cam.py

# 本地视频
python detect_cam.py --source video.mp4

# RTSP 网络摄像头
python detect_cam.py --source rtsp://user:pass@ip:554/stream1

# 换模型 / 调阈值
python detect_cam.py --weights yolov5m.pt --conf-thres 0.4

脚本已包含：

• 实时 FPS 显示
• 按 s 截图
• 按 q 退出

直接跑即可！