YOLOv5保姆级教程:从环境搭建到模型部署,一文搞懂目标检测核心技能
目标检测是计算机视觉领域最基础也最重要的任务之一,从自动驾驶的汽车识别行人,到工业质检中的缺陷检测,再到日常生活中的人脸识别,都离不开目标检测技术的支撑。而在众多目标检测框架中,YOLOv5凭借其出色的性能、简洁的代码结构和丰富的应用场景,成为了工业界和学术界的首选方案。
本文将带你从零开始,系统性地掌握YOLOv5的核心使用方法。无论你是计算机视觉领域的初学者,还是希望将目标检测技术应用到实际项目的工程师,都能从中获得实用的知识和技能。我们会涵盖环境搭建、数据准备、模型训练、推理部署等完整流程,并通过大量实战代码帮助你快速上手。
为什么值得关注:YOLOv5的核心优势解析
YOLOv5是由Ultralytics公司开发和维护的新一代YOLO目标检测算法,相比之前的版本,它在多个方面实现了质的飞跃。首先从性能角度来看,YOLOv5提供了从nano到xlarge五种不同规模的模型架构,能够在精度和速度之间实现灵活平衡。在COCO数据集上的测试结果显示,YOLOv5x达到了50.7%的mAP(平均精度均值),而YOLOv5n在移动端设备上能够实现每秒超过100帧的推理速度,这种灵活的性能选择是其他框架难以比拟的。
从工程实践的角度来看,YOLOv5的代码结构经过精心设计,模块化程度很高。项目的核心代码都集中在models目录下,训练、验证、推理的逻辑清晰易懂,这为二次开发和定制提供了极大的便利。同时,YOLOv5原生支持ONNX、TensorRT、CoreML等多种模型格式,能够轻松部署到边缘设备、云端服务器和移动应用中。这种部署灵活性在实际项目中非常重要,因为不同场景对推理速度和模型大小有着截然不同的要求。
YOLOv5还提供了完整的训练日志和可视化功能。通过集成的Weights and Biases(W&B)支持,你可以实时监控训练过程中的损失曲线、mAP变化、模型性能对比等关键指标。项目自带的plot_results.py脚本能够自动生成训练结果的可视化图表,帮助你快速评估模型效果。此外,YOLOv5的社区非常活跃,GitHub仓库积累了大量的问题解决方案和最佳实践,遇到问题时很容易找到参考。
在实际应用层面,YOLOv5展示了强大的多任务处理能力。除了标准的目标检测任务外,它还支持实例分割(通过YOLOv5的seg分支)、分类任务,以及与DeepSort等跟踪算法结合实现多目标跟踪。这些扩展能力使得YOLOv5能够应对从简单到复杂的各种视觉任务,真正成为一个多功能的视觉基础模型。
环境搭建:完整的开发环境配置指南
在开始使用YOLOv5之前,我们需要先配置好开发环境。一个合理的环境配置不仅能避免后续的各种兼容性问题,还能确保模型训练的效率和稳定性。
首先是最基础的Python环境要求。YOLOv5推荐使用Python 3.8或更高版本,建议使用Python 3.9或3.10以获得最佳兼容性。在创建虚拟环境时,可以使用conda或venv两种方式。以conda为例,我们创建一个专门用于YOLOv5的环境:
# 创建并激活YOLOv5专用环境
conda create -n yolov5 python=3.10
conda activate yolov5
GPU环境配置是至关重要的一步。如果你拥有NVIDIA显卡,正确的CUDA和cuDNN配置将直接决定训练速度。首先需要确认你的显卡驱动版本,然后安装对应的CUDA Toolkit。可以通过nvidia-smi命令查看驱动支持的最高CUDA版本:
# 查看显卡信息和驱动支持的CUDA版本
nvidia-smi
大多数现代NVIDIA显卡(如RTX 3060及以上)都支持CUDA 11.8或更高版本。安装PyTorch时需要根据你的CUDA版本选择正确的安装命令:
# 安装PyTorch(以CUDA 11.8为例)
pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118
克隆YOLOv5仓库后,需要安装项目依赖。官方仓库的requirements.txt文件包含了所有必要的依赖包,但需要注意的是,不同版本的YOLOv5可能对依赖版本有特定要求。建议首先查看当前仓库README中推荐的依赖版本:
# 克隆YOLOv5官方仓库
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
# 安装项目依赖
pip install -r requirements.txt
requirements.txt中的核心依赖包括PyTorch、OpenCV、Matplotlib、Pillow等图像处理库,以及用于模型导出的onnx、tensorboard等工具包。在实际安装过程中,可能会遇到一些版本冲突问题,特别是numpy和opencv-python的版本兼容性。遇到这类问题时,可以尝试指定兼容的版本:
# 如果遇到版本冲突,可以尝试指定版本安装
pip install numpy==1.23.5 opencv-python==4.8.0.74
对于需要使用wandb进行训练可视化的用户,还需要额外安装wandb库并进行登录绑定:
# 安装并登录wandb
pip install wandb
wandb login
完成以上步骤后,可以通过运行一个简单的推理测试来验证环境配置是否正确:
# 验证YOLOv5环境
python detect.py --source data/images/zidane.jpg --weights yolov5s.pt
如果一切配置正确,这个命令会使用预训练模型对示例图片进行目标检测,并在runs/detect/目录下生成带有检测框的结果图片。
核心功能:YOLOv5架构深度剖析
理解YOLOv5的核心架构是掌握这个框架的关键。整个项目可以划分为几个主要模块:模型定义、数据处理、训练逻辑、推理流程和辅助工具。每个模块都有其特定的职责,理解它们之间的关系能够帮助你更好地进行定制开发。
模型架构是YOLOv5的核心。官方提供了六种不同规模的模型配置,从最小的YOLOv5n到最大的YOLOv5x6e。这些配置主要在模型的深度(层数)和宽度(通道数)上有所区别。选择合适的模型规模需要综合考虑任务难度、计算资源、推理速度要求等因素。对于资源受限的移动端场景,YOLOv5n或YOLOv5s是更好的选择;而对于需要高精度的服务器端应用,YOLOv5m或YOLOv5l可能更合适。
YOLOv5的骨干网络(Backbone)采用了CSP(Cross Stage Partial)结构,这种设计能够有效减少计算量的同时保持特征提取能力。骨干网络包含多个CBS(Conv + BatchNorm + SiLU)模块和SPPF(Spatial Pyramid Pooling – Fast)模块。SPPF模块通过多个串行的最大池化操作实现多尺度特征融合,相比SPP模块具有更快的计算速度。
颈部网络(Neck)部分采用了PANet(Path Aggregation Network)和FPN(Feature Pyramid Network)的组合结构。这种双向金字塔结构能够有效地融合不同层级的特征图,使得模型能够同时检测不同尺度的目标。对于小目标检测,这种多尺度特征融合尤为重要。
检测头(Head)负责生成最终的检测结果,包括边界框坐标、目标置信度和类别概率。YOLOv5使用锚框(Anchor)机制来预测边界框,预先定义的锚框尺寸能够加速模型的收敛。官方仓库提供了针对COCO数据集优化的锚框值,在自定义数据集上使用时可能需要根据实际数据重新聚类计算。
数据增强是提升模型泛化能力的重要手段。YOLOv5内置了丰富的数据增强策略,包括马赛克增强(Mosaic)、混合增强(MixUp)、随机仿射变换、颜色空间抖动等。其中马赛克增强是最具特色的技术,它将四张不同的图片拼接成一张进行训练,这种方法能够显著增加样本的多样性,提高模型对小目标的检测能力。
实战教程:目标检测完整流程详解
现在我们进入最核心的部分,通过完整的实战案例来学习如何使用YOLOv5完成目标检测任务。整个流程包括数据准备、模型配置、训练优化和推理部署四个主要阶段。
数据准备:构建高质量的目标检测数据集
数据准备是整个目标检测项目中最重要的环节,数据质量直接决定了最终模型的效果。YOLOv5支持YOLO格式的数据集,这种格式简单高效,被广泛采用。
在YOLO格式中,每张图片对应一个同名的txt标注文件,文件内容每一行代表一个目标,格式为“类别索引 中心x坐标 中心y坐标 宽度 高度”。坐标值都是相对于图片尺寸的归一化数值,范围在0到1之间。例如,如果我们有一张包含两个人的图片,类别索引分别为0和1,对应的标注文件内容可能是:
# 标注文件示例:image.txt
# 格式说明:class_id center_x center_y width height
# 类别0表示人,类别1表示车
0 0.352 0.467 0.186 0.534
1 0.723 0.891 0.234 0.178
创建数据集目录结构时,建议采用以下组织方式:
# 数据集目录结构示例
# dataset/
# ├── images/
# │ ├── train/
# │ └── val/
# ├── labels/
# │ ├── train/
# │ └── val/
# └── dataset.yaml
配置文件dataset.yaml是连接数据集和模型的桥梁,它定义了数据集路径、类别数量和类别名称:
# dataset.yaml 配置文件示例
train: dataset/images/train
val: dataset/images/val
# 类别数量
nc: 2
# 类别名称
names: ['person', 'car']
对于数据标注,推荐使用专业的标注工具来提高效率和准确性。LabelImg是一个开源的图像标注工具,支持导出YOLO格式的标注文件:
# 安装LabelImg标注工具
pip install labelImg
# 启动标注界面
labelImg
在使用LabelImg时,首先需要设置保存目录为labels文件夹,然后设置读取图片目录为images文件夹。每次标注前,确保左侧面板的格式选择为YOLO。标注完成后,记得使用快捷键W创建标注框,A/D键切换上下一张图片,Ctrl+S保存当前标注。
模型训练:从配置到优化的完整流程
准备好数据后,就可以开始训练模型了。YOLOv5的训练脚本功能丰富,支持多种训练策略和优化选项。
首先,我们需要确定使用哪种规模的模型作为基础。模型规模的选择应该基于数据集大小、目标类别数和硬件条件。一般而言,数据集越大、目标类别越多,越适合使用较大的模型。如果你的数据集图片数量少于1000张,建议从YOLOv5s开始;如果超过10000张,可以考虑YOLOv5m或更大的模型。
最基本的训练命令只需要指定数据配置和模型权重:
# 基础训练命令
python train.py --data dataset.yaml --weights yolov5s.pt --epochs 300
这个命令会使用YOLOv5s预训练权重作为初始化,在dataset.yaml定义的数据集上训练300个epoch。预训练权重能够大大加快收敛速度,这是因为模型已经学习到了丰富的图像特征,只需要针对新任务进行微调即可。
在实际训练过程中,监控关键指标非常重要。YOLOv5会在每个epoch结束后自动输出训练指标,包括训练损失、验证损失、precision(精确率)、recall(召回率)和mAP等。mAP@0.5表示在IoU阈值为0.5时的平均精度,mAP@0.5:0.95则表示在IoU从0.5到0.95范围内平均精度的均值,后者是更严格的评估指标。
对于更精细的训练控制,我们可以添加更多参数:
# 带完整参数的训练命令
python train.py \
--data dataset.yaml \
--weights yolov5s.pt \
--cfg yolov5s.yaml \
--epochs 300 \
--batch-size 16 \
--img 640 \
--optimizer SGD \
--lr0 0.01 \
--momentum 0.937 \
--weight-decay 0.0005 \
--patience 50 \
--workers 8 \
--device 0
batch-size(批次大小)直接影响训练速度和显存占用。较大的batch-size能够加速训练并使梯度估计更稳定,但会占用更多显存。在选择batch-size时,应该在不过分牺牲训练速度的前提下,尽可能使用较大的值。一个实用的原则是,先将batch-size设为16,然后逐步增加直到GPU显存溢出,最后取溢出前一个稳定的值。
学习率是另一个关键的超参数。YOLOv5默认使用线性学习率热身(Linear Warmup)策略,在训练初期逐渐增加学习率,这有助于模型在早期快速收敛。对于优化器的选择,SGD通常在目标检测任务中表现更好,但如果训练不稳定,可以尝试切换到Adam:
# 使用Adam优化器
python train.py --data dataset.yaml --weights yolov5s.pt --optimizer Adam
早停机制(Early Stopping)能够防止过拟合。当验证集的mAP在连续若干个epoch内没有提升时,训练会自动停止。这个阈值为patience参数,默认值为50。如果你的训练资源有限,可以适当减小这个值:
# 启用早停,设置耐心值为30
python train.py --data dataset.yaml --weights yolov5s.pt --patience 30
训练过程中的数据增强参数也值得关注。YOLOv5使用hyps.yaml文件来管理数据增强策略,文件中定义了各种增强操作的概率和参数。例如,通过调整mosaic、mixup、copy-paste等增强的比例,可以控制训练数据的多样性程度:
# 数据增强配置示例(hyps.yaml片段)
# 各项增强的概率和参数
hsv_h: 0.015 # 色调增强范围
hsv_s: 0.7 # 饱和度增强范围
hsv_v: 0.4 # 亮度增强范围
degrees: 0.0 # 旋转角度范围
translate: 0.1 # 平移比例
scale: 0.5 # 缩放比例
shear: 0.0 # 剪切角度
perspective: 0.0 # 透视变换
flipud: 0.0 # 上下翻转概率
fliplr: 0.5 # 左右翻转概率
mosaic: 1.0 # 马赛克增强概率
mixup: 0.0 # 混合增强概率
推理部署:模型预测与实际应用
训练完成后,接下来就是将模型应用到实际场景中。YOLOv5提供了灵活的推理接口,支持图片、视频、摄像头实时流等多种输入源。
最基本的图片推理只需要指定模型权重和图片路径:
# 图片推理命令
python detect.py --weights runs/train/exp/weights/best.pt --source data/images/bus.jpg
推理结果默认保存在runs/detect/exp/目录下,包含带有检测框标注的图片。如果需要调整检测参数,可以通过命令行指定置信度阈值和IoU阈值:
# 带参数调整的推理命令
python detect.py \
--weights runs/train/exp/weights/best.pt \
--source data/images/ \
--conf 0.5 \
--iou 0.45 \
--img 640 \
--device 0 \
--save-txt \
--save-conf
conf参数控制检测框的置信度阈值,只有置信度高于此值的检测结果才会被保留;iou参数用于非极大值抑制(NMS),用于去除重叠的检测框。save-txt和save-conf参数会让模型额外保存检测结果的文本文件。
在Python代码中调用YOLOv5推理功能也非常简单。通过导入YOLOv5模块,我们可以将其集成到自己的应用程序中:
# Python API 推理示例
import torch
from pathlib import Path
# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt')
# 设置推理参数
model.conf = 0.5 # 置信度阈值
model.iou = 0.45 # IoU阈值
model.classes = [0, 1] # 只检测特定类别
# 单张图片推理
results = model('data/images/bus.jpg')
# 获取结果
results.print() # 打印检测结果
results.save() # 保存带标注的图片
# 获取检测框坐标
boxes = results.xyxy[0] # xyxy格式:[x1, y1, x2, y2, confidence, class]
for box in boxes:
x1, y1, x2, y2, conf, cls = box
print(f"检测到{results.names[int(cls)]},置信度:{conf:.2f}")
对于批量推理场景,可以处理整个文件夹的图片:
# 批量图片推理
import os
from pathlib import Path
# 获取所有待检测图片
image_dir = Path('data/test_images')
image_files = list(image_dir.glob('*.jpg')) + list(image_dir.glob('*.png'))
# 遍历处理每张图片
for img_path in image_files:
results = model(str(img_path))
# 处理检测结果
print(f"处理完成:{img_path.name}")
视频推理是另一个常见的应用场景。YOLOv5可以直接处理视频文件或摄像头输入:
# 视频推理命令
python detect.py --weights runs/train/exp/weights/best.pt --source video.mp4
# 摄像头实时推理(0表示第一个摄像头)
python detect.py --weights runs/train/exp/weights/best.pt --source 0
# RTSP流推理
python detect.py --weights runs/train/exp/weights/best.pt --source rtsp://camera_ip:554/stream
在处理视频时,如果需要实时显示检测结果,可以添加–view-img参数。对于需要保存处理后视频的场景,使用–save-vid参数会自动保存带检测框的视频文件:
# 保存视频推理结果
python detect.py \
--weights runs/train/exp/weights/best.pt \
--source video.mp4 \
--save-vid \
--view-img
常见应用场景:实战案例分析
YOLOv5的强大之处在于其广泛的适用性。下面我们通过几个实际应用场景来展示如何针对不同任务进行定制开发。
工业质检中的缺陷检测
工业质检是目标检测技术的重要应用领域之一。以电子产品表面缺陷检测为例,我们可以使用YOLOv5来自动识别PCB板上的焊接缺陷、元件缺失、划痕等问题。
首先需要准备包含各类缺陷的标注数据集。常见的缺陷类型包括划痕(scratch)、气泡(bubble)、缺损(missing)、污染(contamination)等。建立数据集时要注意正负样本的平衡,避免某一类缺陷数量过少导致模型在该类别上表现不佳。
针对工业场景的特殊需求,可能需要进行一些定制化调整。例如,如果缺陷很小且密集,可以考虑使用更高的输入分辨率(如1280)来提高小目标检测能力;如果检测速度要求极高,可以选择YOLOv5n模型并使用TensorRT进行部署。
# 工业质检场景的推理代码示例
class QualityInspector:
def __init__(self, model_path, confidence=0.7):
self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
self.model.conf = confidence
self.defect_names = ['scratch', 'bubble', 'missing', 'contamination', 'normal']
def inspect(self, image_path):
results = self.model(image_path)
detections = results.xyxy[0]
defect_count = {}
for det in detections:
cls = int(det[5])
defect_type = self.defect_names[cls]
defect_count[defect_type] = defect_count.get(defect_type, 0) + 1
is_passed = 'normal' in defect_count and len(detections) == defect_count['normal']
return {
'passed': is_passed,
'defects': defect_count,
'image': results.render()[0]
}
智慧交通中的车辆分析
交通领域是目标检测的另一大应用场景。YOLOv5可以用于车牌检测、交通标志识别、车辆统计、拥堵分析等任务。在这类应用中,往往需要处理视频流数据并进行实时分析。
# 视频流处理示例
import cv2
from collections import defaultdict
class TrafficAnalyzer:
def __init__(self, model_path):
self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
self.vehicle_classes = [2, 3, 5, 7] # COCO数据集中车辆相关类别
self.class_names = self.model.names
self.counts = defaultdict(int)
def process_video(self, video_path, output_path=None):
cap = cv2.VideoCapture(video_path)
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
if output_path:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
frame_id = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# YOLOv5推理
results = self.model(frame)
detections = results.xyxy[0]
# 统计车辆数量
for det in detections:
cls = int(det[5])
if cls in self.vehicle_classes:
self.counts[self.class_names[cls]] += 1
# 在视频上标注检测结果
annotated_frame = results.render()[0]
cv2.putText(annotated_frame, f"Frame: {frame_id}", (10, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
if output_path:
out.write(annotated_frame)
frame_id += 1
print(f"处理帧 {frame_id}")
cap.release()
if output_path:
out.release()
return self.counts
智能安防中的人脸与行为检测
在安防领域,YOLOv5可以用于人员检测、入侵报警、行为分析等场景。虽然YOLOv5不是专门的人脸检测模型,但通过训练特定的人体检测模型,可以实现区域入侵检测、陌生人识别等功能。
# 安防监控场景示例
class SecurityMonitor:
def __init__(self, model_path, alert_confidence=0.8):
self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
self.alert_confidence = alert_confidence
self.alert_callback = None
def set_alert_callback(self, callback):
"""设置告警回调函数"""
self.alert_callback = callback
def process_frame(self, frame, draw=True):
results = self.model(frame)
detections = results.xyxy[0]
alerts = []
for det in detections:
conf = det[4].item()
if conf >= self.alert_confidence:
x1, y1, x2, y2 = [int(x) for x in det[:4]]
alerts.append({
'box': (x1, y1, x2, y2),
'confidence': conf
})
if draw:
# 绘制检测框
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
label = f"Person: {conf:.2f}"
cv2.putText(frame, label, (x1, y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
# 触发告警
if alerts and self.alert_callback:
self.alert_callback(alerts)
return frame, alerts
进阶技巧与最佳实践
掌握了基础使用后,了解一些进阶技巧能够帮助你在实际项目中取得更好的效果。
模型性能优化技巧
在实际部署中,推理速度往往是关键考量因素。以下是几种常用的模型优化方法:
使用ONNX格式导出模型能够显著提升推理速度。ONNX是一种通用的模型格式,支持多种推理引擎的优化:
# 导出为ONNX格式
python export.py --weights runs/train/exp/weights/best.pt --include onnx
TensorRT优化能够进一步提升NVIDIA GPU上的推理速度:
# 导出为TensorRT格式
python export.py --weights runs/train/exp/weights/best.pt --include engine
量化技术可以减小模型体积并加快推理速度。INT8量化是一种常用的高效量化方式:
# INT8量化导出(需要校准数据)
python export.py \
--weights runs/train/exp/weights/best.pt \
--include engine \
--int8 \
--data data/coco128.yaml
训练调优经验
模型训练过程中的一些技巧能够帮助获得更好的效果。多尺度训练是提升模型对不同尺寸目标检测能力的有效方法:
# 启用多尺度训练
python train.py --data dataset.yaml --weights yOLOv5s.pt --multi-scale
冻结骨干网络进行微调可以防止过拟合,特别是当你的数据集较小时:
# 冻结骨干网络前24层
python train.py --data dataset.yaml --weights yolov5s.pt --freeze 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
针对类别不平衡问题,可以使用类别权重调整:
# 使用加权损失(需要在训练代码中添加此功能)
# 在utils/loss.py中添加类别权重计算逻辑
数据集质量提升策略
高质量的数据集是获得优秀模型的前提。在标注完成后,进行数据清洗和检查是必要的步骤:
# 检查数据集完整性的脚本
from pathlib import Path
def verify_dataset(data_yaml):
import yaml
with open(data_yaml, 'r') as f:
data = yaml.safe_load(f)
train_images = Path(data['train'])
train_labels = Path(data['train'].replace('images', 'labels'))
missing_labels = []
for img_path in train_images.glob('*.jpg'):
label_path = train_labels / (img_path.stem + '.txt')
if not label_path.exists():
missing_labels.append(img_path.name)
print(f"缺失标注的图片数量:{len(missing_labels)}")
if missing_labels:
print("缺失标注的图片:", missing_labels[:10])
使用数据增强来扩充小样本类别是另一个重要策略:
# 自定义数据增强脚本
import albumentations as A
from PIL import Image
import cv2
def augment_image(image_path, label_path, output_dir, num_augments=5):
image = cv2.imread(str(image_path))
with open(label_path, 'r') as f:
labels = f.readlines()
transform = A.Compose([
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.5),
A.RandomGamma(p=0.3),
A.Blur(blur_limit=3, p=0.3),
A.MotionBlur(blur_limit=3, p=0.2),
], bbox_params=A.BboxParams(format='yolo'))
for i in range(num_augments):
augmented = transform(image=image, bboxes=[parse_yolo_label(l) for l in labels])
aug_image = augmented['image']
aug_labels = augmented['bboxes']
# 保存增强后的图片和标签
cv2.imwrite(f"{output_dir}/aug_{i}_{image_path.name}", aug_image)
save_yolo_labels(f"{output_dir}/aug_{i}_{label_path.name}", aug_labels)
模型部署实战
将训练好的模型部署到生产环境是项目的最终目标。YOLOv5支持多种部署方式,下面介绍几种常见的部署方案。
Flask Web服务部署
将YOLOv5封装成Web API是常见的部署方式之一,便于与各种前端应用集成:
# Flask API服务示例
from flask import Flask, request, jsonify
import torch
from PIL import Image
import io
app = Flask(__name__)
# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom',
path='runs/train/exp/weights/best.pt')
@app.route('/detect', methods=['POST'])
def detect():
if 'image' not in request.files:
return jsonify({'error': 'No image provided'}), 400
image_file = request.files['image']
image_bytes = image_file.read()
image = Image.open(io.BytesIO(image_bytes))
# 执行检测
results = model(image)
# 提取检测结果
detections = []
for *box, conf, cls in results.xyxy[0]:
detections.append({
'class': model.names[int(cls)],
'confidence': float(conf),
'bbox': [float(x) for x in box]
})
return jsonify({
'success': True,
'detections': detections,
'count': len(detections)
})
@app.route('/health', methods=['GET'])
def health():
return jsonify({'status': 'healthy'})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)
Docker容器部署
使用Docker能够确保模型在不同环境中的一致性:
# Dockerfile
FROM python:3.10-slim
WORKDIR /app
# 安装系统依赖
RUN apt-get update && apt-get install -y \
libgl1-mesa-glx \
libglib2.0-0 \
&& rm -rf /var/lib/apt/lists/*
# 复制requirements并安装Python依赖
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# 复制模型文件
COPY best.pt /app/weights/best.pt
# 复制应用代码
COPY app.py .
# 暴露端口
EXPOSE 5000
# 运行应用
CMD ["python", "app.py"]
# docker-compose.yml
version: '3.8'
services:
yolov5-api:
build: .
ports:
- "5000:5000"
environment:
- PYTHONUNBUFFERED=1
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
边缘设备部署
对于树莓派等边缘设备,可以使用轻量级模型配合优化后的推理库:
# 边缘设备部署优化脚本
import torch
import numpy as np
# 使用TorchScript优化模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5n.pt')
model.eval()
# TorchScript导出
example_input = torch.zeros(1, 3, 640, 640)
traced_model = torch.jit.trace(model.model, example_input)
traced_model.save('yolov5n_traced.pt')
# 量化模型
quantized_model = torch.quantization.quantize_dynamic(
model.model, {torch.nn.Linear}, dtype=torch.qint8
)
总结与资源推荐
通过本文的系统学习,你应该已经掌握了YOLOv5目标检测框架的核心使用方法。从环境搭建到模型训练,从推理部署到性能优化,我们涵盖了完整的目标检测项目开发流程。
YOLOv5作为目前最流行的开源目标检测框架之一,其优势不仅体现在出色的性能上,更重要的是它提供了完整的技术生态和活跃的社区支持。掌握好这个框架,将为你进入计算机视觉领域打下坚实的基础。
在实际工作中,建议按照以下顺序进一步深入学习:首先精通YOLOv5的标准用法,然后学习模型优化和部署技术,接着可以尝试扩展到其他视觉任务如分割、跟踪等,最后可以根据项目需求进行定制化开发。
相关项目推荐
以下是一些与YOLOv5互补的优秀开源项目,可以帮助你构建更完整的计算机视觉解决方案:
YOLOv8是YOLOv5的升级版本,引入了新的架构设计和更强大的性能。如果你需要更高的检测精度,可以尝试使用YOLOv8。仓库地址:github.com/ultralytics/ultralytics
DeepSORT是多目标跟踪算法,与YOLOv5结合可以实现目标的持续跟踪。在交通监控、人流统计等场景中非常有用。仓库地址:github.com/nwojke/deep_sort
SAM(Segment Anything Model)是Meta推出的革命性分割模型,可以实现任意目标的分割。与YOLOv5的目标检测结合,可以构建完整的目标检测和分割pipeline。仓库地址:github.com/facebookresearch/segment-anything
MMDetection是商汤科技开源的目标检测工具箱,支持多种检测算法和丰富的模型库。如果你需要尝试更多检测算法,MMDetection是很好的选择。仓库地址:github.com/open-mmlab/mmdetection
最后,建议持续关注YOLOv5官方仓库的更新,开发团队会不断发布新功能和性能优化。同时,参与社区讨论和贡献代码也是提升技术水平的好方法。祝你在使用YOLOv5的过程中收获满满!
评论区