4. YOLOv5推理与图像目标检测

任务目标

加载一个训练好的YOLOv5模型，对单张输入图片执行前向推理，解析输出得到每个目标的类别、置信度与边界框坐标，并在原图上绘制可视化结果后保存。本案例聚焦推理闭环，不涉及训练过程。

环境与依赖

需要Python 3.8+，并安装以下依赖（建议使用conda或venv隔离环境）：

pip install torch torchvision opencv-python numpy

YOLOv5代码库通过Git获取（推理时仅用到其模型定义与权重加载逻辑）：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

关键完整代码

以下代码完成从加载模型到保存可视化结果的完整推理流程。将代码保存为infer.py，放在yolov5目录下运行。

import cv2
import torch
import numpy as np

# 加载预训练模型（自动下载COCO权重）
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
model.conf = 0.4    # 置信度阈值
model.iou = 0.45    # NMS IoU阈值

# 读取输入图片
img_path = 'bus.jpg'   # 替换为你的图片路径
img = cv2.imread(img_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 执行推理
results = model(img_rgb)

# 解析结果：获取边界框、置信度、类别ID
boxes = results.xyxy[0].cpu().numpy()  # [x1, y1, x2, y2, conf, class]

# 在图片上绘制检测框与标签
for det in boxes:
    x1, y1, x2, y2, conf, cls_id = det.astype(float)
    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    cls_id = int(cls_id)
    label = f"{model.names[cls_id]} {conf:.2f}"

    # 绘制矩形框
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    # 绘制标签背景
    (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
    cv2.rectangle(img, (x1, y1 - h - 10), (x1 + w, y1), (0, 255, 0), -1)
    # 绘制标签文字
    cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)

# 保存结果
output_path = 'output.jpg'
cv2.imwrite(output_path, img)
print(f"检测到 {len(boxes)} 个目标，结果已保存至 {output_path}")

运行方式

1. 准备一张包含常见物体（如人、车、动物）的测试图片，命名为bus.jpg（或修改代码中的img_path变量）。
2. 在yolov5目录下执行：python infer.py
3. 首次运行会自动下载yolov5s.pt权重（约14MB），需保持网络畅通。

验证方法

程序正常结束且无报错后，检查以下两点判定推理成功：

• 控制台输出检测到的目标数量，且数量大于0（若图片中确有可检测物体）。
• 打开生成的output.jpg，确认图中目标被绿色矩形框标出，框上方显示类别名称与置信度数值，且框位置与目标实际位置基本吻合。

若输出图片为原图无任何框，可尝试降低model.conf至0.25后重新运行；若仍无检测结果，说明图片中物体不在COCO 80类范围内或目标过小。

常见问题

• CUDA不可用：代码自动回退到CPU推理，速度较慢但结果正确。如需GPU加速，确保安装对应版本的PyTorch。
• 权重下载失败：可手动下载yolov5s.pt放入当前目录，代码会自动识别本地权重。
• 中文路径问题：确保图片路径不含中文，否则OpenCV读取可能失败。