YOLO-v5代码实例:加载自定义图片进行推理的详细步骤
YOLO(You Only Look Once)是一种流行的物体检测和图像分割模型,由华盛顿大学的Joseph Redmon 和Ali Farhadi 开发。 YOLO 于2015 年推出,因其高速和高精度而广受欢迎。经过多个版本迭代,YOLOv5 在保持实时性的同时进一步提升了检测精度与易用性,成为工业界和学术界广泛应用的目标检测解决方案之一。
基于YOLOv5算法构建的深度学习镜像,提供了完整的计算机视觉开发环境。该镜像预装了PyTorch深度学习框架及YOLOv5专用工具库,支持目标检测任务的快速开发与部署。本文将详细介绍如何使用该镜像加载自定义本地图片并执行推理,涵盖环境准备、代码实现、结果处理等关键环节,帮助开发者快速上手YOLOv5的实际应用。
1. 环境与资源准备
在开始推理之前,确保已正确配置YOLOv5运行环境。本节介绍两种常见使用方式:Jupyter Notebook 和 SSH 远程访问,并说明项目目录结构。
1.1 使用 Jupyter Notebook
Jupyter 是交互式开发的理想选择,特别适合调试和可视化。启动容器后,可通过浏览器访问 Jupyter 页面:
- 打开浏览器,输入提供的 URL 地址
- 导航至
/root/yolov5/目录 - 创建新的
.ipynb文件或打开已有示例 notebook
1.2 使用 SSH 方式连接
对于需要长期运行或批量处理的任务,推荐通过 SSH 登录服务器进行操作:
- 使用终端执行
ssh root@<your-server-ip>进行连接 - 输入密码后进入系统
- 可直接运行 Python 脚本或启动后台任务
1.3 项目路径与依赖确认
首先进入 YOLOv5 主目录:
cd /root/yolov5/该目录包含以下核心组件:
models/:模型定义文件(如 yolov5s.yaml)utils/:通用工具函数(绘图、数据增强等)detect.py:标准推理脚本requirements.txt:所需 Python 包列表
所有依赖均已预安装,无需额外配置即可调用torch.hub.load加载模型。
2. 模型加载与图像输入设置
本节将演示如何从 Hub 加载预训练 YOLOv5 模型,并准备自定义图像作为输入源。
2.1 加载 YOLOv5 预训练模型
YOLOv5 提供多种尺寸模型以适应不同硬件条件和性能需求:
| 模型 | 参数量(M) | 推理速度(FPS) | 适用场景 |
|---|---|---|---|
| yolov5n | 1.9 | ~160 | 边缘设备、低延迟要求 |
| yolov5s | 7.2 | ~90 | 平衡精度与速度 |
| yolov5m | 21.2 | ~50 | 中等算力平台 |
| yolov5l | 46.5 | ~30 | 高精度需求 |
| yolov5x | 86.7 | ~18 | 最高精度 |
默认推荐使用yolov5s,兼顾效率与准确性:
import torch # Load a YOLOv5 model (options: yolov5n, yolov5s, yolov5m, yolov5l, yolov5x) model = torch.hub.load("ultralytics/yolov5", "yolov5s") # Default: yolov5s首次运行时会自动下载权重文件(约 140MB),后续调用将直接从缓存加载。
2.2 设置自定义图像输入
YOLOv5 支持多种图像输入格式,包括:
- 网络图片 URL
- 本地文件路径(
.jpg,.png等) - PIL 图像对象
- OpenCV 的 numpy 数组(BGR 格式)
- 多图列表(用于批量推理)
示例:加载本地图片
假设你有一张名为custom_image.jpg的图片上传至/root/yolov5/data/images/目录下:
img = "/root/yolov5/data/images/custom_image.jpg"若目录不存在可手动创建:
mkdir -p /root/yolov5/data/images # 将图片上传至此目录也可以使用 OpenCV 读取图像并传入模型:
import cv2 img_cv = cv2.imread("/root/yolov5/data/images/custom_image.jpg") results = model(img_cv) # 自动转换 BGR → RGB3. 执行推理与结果解析
模型加载完成后,即可对输入图像执行端到端推理。YOLOv5 内部自动完成图像预处理(归一化、缩放、填充),用户无需手动干预。
3.1 执行推理并输出结果
# Perform inference results = model(img) # Print detection results to console results.print()输出示例如下:
image 1/1 /root/yolov5/data/images/custom_image.jpg: 640x480 4 persons, 1 tie, Done. (0.023s)表示检测到 4 个人和 1 条领带,推理耗时 23ms。
3.2 可视化检测结果
YOLOv5 提供便捷的可视化方法:
results.show() # 弹窗显示带框图像(需 GUI 环境)注意:在无图形界面的服务器上运行时,
show()可能报错。建议改用save()保存图像。
results.save() # 保存结果到 runs/detect/exp/保存路径按时间戳递增命名(如exp,exp2),便于区分多次实验。
3.3 获取结构化检测数据
除了可视化,还可提取结构化结果用于后续分析:
# Convert to pandas DataFrame df = results.pandas().xyxy[0] # 第一张图的结果 print(df)输出字段说明:
| 列名 | 含义 |
|---|---|
| xmin, ymin | 检测框左上角坐标 |
| xmax, ymax | 检测框右下角坐标 |
| confidence | 置信度分数(0~1) |
| name | 类别名称(如 'person', 'car') |
| class | 类别索引(COCO 数据集共80类) |
可用于过滤特定类别或高置信度结果:
# 提取所有人脸(假设类别为 'person' 且置信度 > 0.7) high_conf_persons = df[(df['name'] == 'person') & (df['confidence'] > 0.7)]3.4 裁剪检测区域(Crop)
YOLOv5 支持一键裁剪出每个检测对象:
cropped_images = results.crop(save=True) # 保存裁剪后的图像裁剪结果保存在runs/detect/exp/crops/<class_name>/目录下,适用于人脸识别、商品识别等子任务。
4. 完整代码示例与最佳实践
以下是整合上述步骤的完整可运行代码模板,适用于大多数自定义图像推理场景。
4.1 完整推理脚本
import torch import os from pathlib import Path # --- Step 1: Load Model --- model = torch.hub.load("ultralytics/yolov5", "yolov5s") # --- Step 2: Define Image Source --- img_path = "/root/yolov5/data/images/custom_image.jpg" # Check if file exists if not os.path.exists(img_path): raise FileNotFoundError(f"Image not found at {img_path}") # --- Step 3: Inference --- results = model(img_path) # --- Step 4: Process Results --- # Print to console results.print() # Save annotated image results.save() # Export cropped objects results.crop(save=True) # Get structured data df = results.pandas().xyxy[0] print("\nDetection Details:") print(df[['name', 'confidence', 'xmin', 'ymin', 'xmax', 'ymax']])4.2 常见问题与解决方案
| 问题 | 原因 | 解决方案 |
|---|---|---|
ModuleNotFoundError: No module named 'yolov5' | 路径未加入 PYTHONPATH | 运行前执行export PYTHONPATH="$PYTHONPATH:/root/yolov5" |
cv2.imshow()报错 | 无 GUI 环境 | 改用results.save()替代show() |
| 推理结果为空 | 图像模糊或目标过小 | 尝试使用更大模型(如 yolov5m)或调整输入分辨率 |
| 下载模型失败 | 网络超时 | 手动下载权重并指定本地路径加载 |
4.3 性能优化建议
- 启用半精度(FP16)加速:
model = model.half().cuda() # GPU 上启用 float16 img_tensor = torch.half().cuda() # 输入也转为 half- 调整图像大小以提升速度:
results = model(img, size=320) # 默认 640,降低可提速但影响精度- 批量推理提高吞吐量:
img_list = ["img1.jpg", "img2.jpg", "img3.jpg"] results = model(img_list)5. 总结
本文详细介绍了如何在 YOLOv5 深度学习镜像环境中加载自定义图片并执行目标检测推理。我们从环境搭建入手,逐步讲解了模型加载、图像输入设置、推理执行、结果解析与可视化等全流程操作,并提供了完整的可运行代码示例。
通过本教程,读者可以掌握以下核心技能:
- 正确使用
torch.hub.load加载不同规模的 YOLOv5 模型; - 灵活传入本地图片、OpenCV 图像等多种输入形式;
- 利用
.print(),.save(),.crop()等方法高效处理检测结果; - 提取结构化数据(pandas DataFrame)用于下游分析;
- 应对常见问题并进行性能优化。
YOLOv5 凭借其简洁的 API 设计和强大的功能集成,极大降低了目标检测技术的应用门槛。结合预置镜像环境,开发者可在几分钟内完成模型部署与测试,加速 AI 项目的落地进程。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。