3、YOLOv5模型加载：使用PyTorch Hub加载YOLOv5s模型、模型推理、结果解析

模型加载这一步，说白了就是让我们的程序「认识」YOLOv5。我个人习惯用 PyTorch Hub 来加载，因为它最省事——不用手动下载权重文件，也不用纠结环境依赖。你想想看，一个 torch.hub.load() 调用，模型和权重就全到位了。

3.1 为什么选 PyTorch Hub？

我记得最早做目标检测时，还得自己去 GitHub 上 clone 仓库、装 requirements、下载 .pt 文件……折腾半天。后来 PyTorch Hub 普及了，一切都简单了。

PyTorch Hub 的好处就三点：

• 自动下载：模型结构和权重一次搞定
• 版本隔离：不同项目可以用不同版本的模型
• 缓存机制：第二次加载秒开

3.2 加载 YOLOv5s 模型

YOLOv5 有多个版本：n、s、m、l、x。s 代表 small，是速度和精度的平衡点。我在项目中做过测试，v5s 在 1080Ti 上能跑到 100+ FPS，精度也够用。

加载代码其实就一行：

import torch

# 加载 YOLOv5s 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

嗯，这里要注意：第一次加载会下载模型文件（大约 14MB），网络慢的话可能要等一会儿。加载成功后，你会看到类似这样的输出：

Downloading: "https://github.com/ultralytics/yolov5/zipball/master" to /root/.cache/torch/hub/master.zip
...
YOLOv5 🚀 2023-6-26 Python-3.9.16 torch-2.0.1 CUDA:0 (Tesla T4, 15102MiB)

看到这个就说明模型加载成功了。我个人习惯把模型放到 GPU 上：

# 如果有 GPU 就使用 GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
print(f'模型已加载到 {device}')

3.3 模型推理

模型加载好了，接下来就是推理。YOLOv5 的推理接口非常友好，支持多种输入格式：

• 单张图片路径（字符串）
• PIL Image 对象
• OpenCV 读取的 numpy 数组
• 图片列表

我最常用的是 numpy 数组，因为摄像头读取的就是这个格式：

import cv2

# 读取图片
img = cv2.imread('test.jpg')
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 推理
results = model(img_rgb)

# 看看结果
results.print()  # 打印检测结果
results.show()   # 显示检测后的图片

你可能会问：为什么要转成 RGB？因为 YOLOv5 训练时用的是 RGB 格式，而 OpenCV 默认是 BGR。不转换的话，颜色会偏蓝，检测效果也会受影响。

推理速度方面，我实测过：

3.4 结果解析

推理完成后，results 对象里包含了所有检测信息。我刚开始用的时候，直接 results.pandas().xyxy[0] 一把梭，后来发现这样不够灵活。咱们来看看 results 到底长什么样：

# 查看原始结果
print(results.xyxy[0])  # tensor 格式
print(results.pandas().xyxy[0])  # pandas DataFrame 格式

输出示例：

          xmin        ymin        xmax        ymax  confidence  class  name
0   7.531250e+01  233.500000  114.750000  316.500000    0.894531      0  person
1   1.148438e+02  195.250000  229.500000  410.500000    0.871094      0  person
2   4.750000e+02   73.500000  549.000000  194.500000    0.835938     39  bottle

每一行代表一个检测到的目标，包含：

• xmin, ymin, xmax, ymax：边界框坐标（左上角和右下角）
• confidence：置信度，0~1 之间
• class：类别编号
• name：类别名称

在实际项目中，我通常会这样解析：

# 获取检测结果
detections = results.xyxy[0].cpu().numpy()

for det in detections:
    x1, y1, x2, y2, conf, cls = det
    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    label = results.names[int(cls)]
    
    print(f'检测到 {label}，置信度 {conf:.2f}，位置 ({x1},{y1})-({x2},{y2})')
    
    # 过滤低置信度检测
    if conf < 0.5:
        continue
    
    # 这里可以画框、计数、做后续处理

3.5 完整流程串联

把上面这些串起来，就是一个完整的检测流程。我画了张图，帮你理清思路：

这个流程看着简单，但实际项目中坑不少。我遇到过最典型的问题就是——模型加载成功了，推理也跑了，结果全是空检测。后来发现是输入图片的尺寸不对，YOLOv5 默认要求 640×640，而我传了 1920×1080 的图。

3.6 完整代码示例

最后，给你一个可以直接跑的完整示例。我在项目中就是这么用的：

import torch
import cv2
import numpy as np

class YOLOv5Detector:
    def __init__(self, conf_thres=0.5):
        # 加载模型
        self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
        self.model.conf = conf_thres  # 置信度阈值
        self.model.iou = 0.45         # NMS 阈值
        
        # 使用 GPU
        if torch.cuda.is_available():
            self.model.to('cuda')
            print('使用 GPU 推理')
        else:
            print('使用 CPU 推理')
    
    def detect(self, img):
        # img: numpy array, shape (H, W, 3), BGR 格式
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        results = self.model(img_rgb)
        return results
    
    def parse_results(self, results):
        detections = []
        for det in results.xyxy[0].cpu().numpy():
            x1, y1, x2, y2, conf, cls = det
            detections.append({
                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                'confidence': float(conf),
                'class_id': int(cls),
                'class_name': results.names[int(cls)]
            })
        return detections

# 使用示例
detector = YOLOv5Detector(conf_thres=0.5)
img = cv2.imread('test.jpg')
results = detector.detect(img)
detections = detector.parse_results(results)

for d in detections:
    print(f"{d['class_name']}: {d['confidence']:.2f}")