YOLOv5垃圾分类检测系统实战：从数据集训练到PyQt5 UI界面集成

最近把智能生活垃圾检测与分类系统完整做了一遍从数据集整理、YOLOv5模型训练到UI界面集成一路跑通。这套东西说难不难但细节相当多尤其是“UI界面卡顿”和“训练自己的数据集”这两块几乎是每个做类似项目的人都会撞上的墙。这篇文章把我从0到1的实现过程、踩过的坑、调通的参数全部写出来给正在折腾同类项目的朋友一个可以照抄的参考。1. 系统整体设计与技术选型思路1.1 项目定位与核心需求这套系统要解决的事情很直白用户把垃圾放到摄像头前或者上传一张图片系统自动识别这是什么垃圾属于哪个分类并把结果实时显示在界面上。应用场景可以是一个智能回收站的识别终端也可以是小区垃圾分类亭的辅助判断设备甚至能做成手机端的小程序。我给自己定的技术指标是单帧检测延迟不超过150ms普通摄像头场景够用UI界面操作流畅不卡顿支持图片和视频流两种输入方式检测结果可统计导出。整体架构分成三层数据层负责图像采集与输入算法层跑YOLOv5做目标检测与分类展示层用UI界面呈现检测框、类别标签和置信度。1.2 为什么是YOLOv5而不是其他模型模型选型我对比了好几个方向YOLOv5、YOLOv8、SSD-MobileNet、Faster R-CNN。最终选YOLOv5不是因为它是参数最漂亮的那个而是从工程落地角度它最稳。模型推理速度训练难度社区资源部署生态综合判断YOLOv5s快GPU实时低极其丰富成熟支持ONNX/TensorRT首选YOLOv8n更快低丰富成熟备选新项目可用SSD-MobileNet快中一般一般精度偏低不推荐Faster R-CNN慢高一般较复杂精度高但实时性差YOLOv5的社区生态是我最看重的。训练自己的数据集YOLOv5的教程数量、踩坑案例、预训练权重都是这几个方案里最多的遇到问题几乎都能搜到答案。而且它的代码结构清晰detect.py、train.py、data.yaml这些关键文件分工明确对刚接触目标检测项目的人非常友好。如果你之后再迁移到YOLOv8很多思路也是通用的。1.3 UI界面方案的选择UI方案我纠结过几条路PyQt5桌面程序、Tkinter轻量界面、Web前端FlaskHTML页面。考虑到项目要求是“UI界面YOLOv5”我最终选了PyQt5。主要原因有三个第一PyQt5做桌面端视频实时显示有天然优势信号槽机制和QThread多线程配合能解决“UI界面卡顿”这个核心问题第二PyQt5打包成exe方便演示和部署教室或社区场景一台电脑就能跑第三PyQt5有QSS样式表能用简单代码做出好看UI不需要额外写前端。Web方案不是不行但Flask浏览器方案在处理视频流实时推送时要额外解决WebSocket转发或MJPEG流封装的问题复杂度反而更高。Tkinter则太简陋做出来的界面观感一般不适合作为项目展示。2. 训练数据集的构建决定模型上限的核心环节2.1 数据来源与公开数据集数据是这套系统的基础训练集质量直接决定模型识别准不准。我用的数据主要来自两个公开数据集Kaggle上的Garbage Classification数据集包含纸板、玻璃、金属、纸张、塑料、厨余垃圾等6个常见类别以及华为云垃圾分类比赛的数据集类别更多包含40多个小类但有部分图片质量参差不齐。我的具体做法是先用公开数据集搭建一个能跑通的基础版本再手动补充自采数据来改善特定类别的识别能力。比如“塑料瓶”这一类公开数据集里的图片偏少且背景单一我就在办公室和校园里拍了大概200张不同角度、不同光照的塑料瓶照片补进去。实测下来补充自采数据后这类别的mAP提升非常明显。提示不要盲目追求类别多。垃圾分类系统建议控制在8类以内类别越细边界越模糊误判率越高。比如“纸盒”和“纸张”在形态上差异很大但“纸杯”和“纸盒”就容易混淆需要额外注意数据标注的边界规范。2.2 数据清洗与标注规范拿到原始图片后第一步不是急着标注而是做数据清洗。我筛掉了三种图片分辨率低于300x300的、严重模糊或过曝的、带大量水印的。清洗比例大概占原始数据的5%这个步骤很多人会跳过但保留低质量图只会给训练增加噪声。标注工具我用的是LabelImg操作简单、支持VOC格式导出、跨平台。标注时我给自己定了几条规范目标物体占画面比例低于10%的图片直接弃用小目标会让模型学习困难同一物体必须完整框住不截半遮挡严重的物体跳过不标类别标签用英文如plastic_bottle、paper_box避免中文路径和中文标签带来的编码问题LabelImg操作很直观左侧工具栏选“Create RectBox”在图片上拖出矩形框然后选类别按CtrlS保存。不过要注意LabelImg默认保存的是Pascal VOC的XML格式YOLOv5训练需要的是YOLO格式的txt这一步要做格式转换。2.3 VOC格式转YOLO格式YOLO格式的标签是每行一个目标内容是类别ID cx cy w h其中cx cy是中心点坐标w h是宽高全部除以图片宽高归一化到0-1之间。手动算太麻烦我用脚本批量处理。import xml.etree.ElementTree as ET import os from PIL import Image def convert_voc_to_yolo(xml_path, img_dir, save_dir, class_names): tree ET.parse(xml_path) root tree.getroot() img_name root.find(filename).text img_path os.path.join(img_dir, img_name) img_w, img_h Image.open(img_path).size txt_content [] for obj in root.iter(object): class_name obj.find(name).text class_id class_names.index(class_name) bbox obj.find(bndbox) xmin float(bbox.find(xmin).text) ymin float(bbox.find(ymin).text) xmax float(bbox.find(xmax).text) ymax float(bbox.find(ymax).text) x_center (xmin xmax) / 2 / img_w y_center (ymin ymax) / 2 / img_h w (xmax - xmin) / img_w h (ymax - ymin) / img_h txt_content.append(f{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) txt_name os.path.splitext(img_name)[0] .txt with open(os.path.join(save_dir, txt_name), w) as f: f.write(\n.join(txt_content)) class_names [cardboard, glass, metal, paper, plastic, trash] convert_voc_to_yolo(data/annotations/img001.xml, data/images, data/labels, class_names)转换时有个容易踩的坑XML里的filename字段和实际文件名不一致。有些数据集改过名字但XML里还是旧名称导致脚本找不到图片。我的做法是先解析一个样本打印出filename和实际文件列表对比一下确认一致后再批量处理。2.4 数据划分与增强策略数据做训练集和验证集的划分比例我用的9:1比常用的8:2略高一点。因为垃圾分类数据相对不复杂验证集主要用来监控过拟合趋势和mAP指标不需要留特别多。目录结构我严格遵守YOLOv5的约定dataset/ ├── images/ │ ├── train/ │ └── val/ ├── labels/ │ ├── train/ │ └── val/ └── data.yaml图片和标签的文件名要一一对应比如img001.jpg对应img001.txt。这个对应关系少了后缀的部分YOLOv5是通过遍历图片目录再找同名txt的方式建立的所以两个目录的文件名不含扩展名必须完全一致。数据增强主要靠YOLOv5内置策略Mosaic、随机翻转、HSV色域增强等。我不是很建议再额外做离线增强比如旋转、加噪因为YOLOv5训练时默认就会做在线增强再加一轮反而容易破坏物体真实比例特征。尤其是Mosaic增强它把4张图拼成一张训练对于小物体检测很有帮助但训练轮数过多时可能引起过拟合需要配合早停机制控制。最终我的数据集规模是约8000张图片每类1000-1500张类别按可回收物场景分了6个纸板、玻璃、金属、纸张、塑料瓶、厨余垃圾。3. YOLOv5训练自己的数据集环境、配置与坑点3.1 环境配置版本匹配比新版本更重要YOLOv5的训练环境配置我这里直接给出经过验证的稳定组合Python 3.8 PyTorch 1.10.0 CUDA 11.3 cuDNN 8.2.0。这套组合稳定性和兼容性很成熟网上资料也多遇到问题好解决。注意很多人一上来就装最新版Python和PyTorch结果遇到torch.cuda.is_available()返回False。这多半是CUDA版本和PyTorch版本不匹配或者安装成了CPU版。先确定自己的NVIDIA驱动支持哪个CUDA版本再选对应的PyTorch版本顺序不要反。安装依赖时用官方给出的requirements.txt即可git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt如果是在国内网络环境下pip下载PyTorch可能很慢建议用清华或阿里云镜像源。但注意PyTorch官方给的CUDA版本需要从官方源下载镜像源里有些是CPU版安装前看清楚torch.cuda.is_available()能不能返回True。3.2 data.yaml配置一个字母都不能错data.yaml是连接数据集和训练脚本的桥梁里面定义了类别数量和类别名称。train: ../dataset/images/train val: ../dataset/images/val nc: 6 names: [cardboard, glass, metal, paper, plastic, trash]这里最容易出问题的就是names列表的顺序。YOLOv5按列表顺序分配类别ID比如cardboard是0plastic是4。如果你的标签文件里某个txt写的是4对应到YOLOv5里就应该找names[4]即plastic。如果你把names顺序调了训练出来的模型预测结果就全乱了。3.3 训练超参数显存、batch-size、img-size的取舍训练命令我长这样python train.py --data data.yaml --weights yolov5s.pt --img 640 --batch-size 16 --epochs 150 --device 0几个参数的具体考量img 640YOLOv5s默认输入尺寸就是640x640。对小目标检测有帮助但显存占用也线性增加。如果你的显存只有6G建议用--img 416速度更快牺牲一点精度。batch-size这块最容易被OOM卡住。显存8G的卡跑YOLOv5s的640输入batch-size设置16基本是上限显卡不行就设8或4。其实batch-size小一点没关系配合默认优化器训练效果差距不大。epochs新手容易训练不足或者过拟合。150轮对于这个量级的数据集是比较稳的区间。训练过程中注意观察验证集的mAP指标如果100轮以后mAP几乎不涨了就手动CtrlC停掉没必要硬跑满。workersWindows下经常出现DataLoader卡死问题把--workers 0加上可以解决。训练日志里最需要关注的是mAP0.5和mAP0.5:0.95这两个指标。mAP0.5计算的是IoU阈值0.5下的平均精度工程上用这个值作为评判标准比较实用。我最终训练出来的模型mAP0.5达到0.93左右实际测试下来的误判率可以接受。心得训练完成后要同时看loss曲线和mAP曲线。如果loss一路下降但验证集mAP不涨反跌说明过拟合了需要回到两个方向增加数据量、打开更强的数据增强。如果loss降不下去大概率是数据标注有误比如框没框准、类别标错了检查数据比调参更有效。3.4 模型选择从s到x的体积和精度权衡YOLOv5有n/s/m/l/x五个规格体积和精度依次增加。我的建议是普通PC演示、课程设计用YOLOv5s最合适精度和速度平衡模型文件约14MB需要边缘设备如Jetson Nano部署选YOLOv5n模型只有4MB不到推理框很小追求极致精度、对速度不敏感的离线识别可以上YOLOv5l我用的是YOLOv5s在GTX 1660显卡上单帧推理大概18ms算上图像预处理和显示渲染总体延迟不超过100ms完全满足实时要求。训练好的模型会保存在runs/train/exp*/weights/目录其中best.pt是验证集上mAP最高的权重last.pt是最后一轮的权重。实际部署一定用best.pt因为过拟合时last.pt往往已经飘了。4. UI界面设计与YOLOv5推理集成4.1 UI界面功能拆解与布局设计UI界面我设计了四个核心区域左侧大区是视频显示区负责实时显示摄像头画面或加载图片绘制检测框和标签右上角是检测结果面板按类别统计当前画面的垃圾数量右下角是历史统计区展示检测累计数据和分类占比条形图顶部是控制栏包含“打开摄像头”“识别图片”“停止检测”三个核心按钮界面上的字体、颜色、布局我用QSS样式表做了统一美化整体看起来是深色背景加亮色检测框的风格。这套UI“好看”不靠堆图片而是靠间距统一、色系一致、标注层次分明。4.2 UI界面卡顿的根源主线程被阻塞“UI界面卡顿”是这类项目最常见的问题几乎每个人第一次做都会遇到。根因在于YOLOv5推理是CPU/GPU密集型任务如果你直接把推理逻辑写在UI主线程里每一帧推理期间界面都无法响应鼠标操作和画面刷新看起来就是卡死。解决思路是引入多线程推理放在QThread子线程中UI主线程只负责接收推理结果并刷新画面。两个线程之间用信号槽通信子线程发出sigFrame信号把检测完成的图像帧传给主线程主线程负责把它转成QPixmap显示。import sys from PyQt5.QtCore import QThread, pyqtSignal from PyQt5.QtGui import QImage, QPixmap import cv2 import torch class DetectThread(QThread): sigFrame pyqtSignal(QImage) def __init__(self, model, source0): super().__init__() self.model model self.source source self.running True def run(self): cap cv2.VideoCapture(self.source) while self.running: ret, frame cap.read() if not ret: break results self.model(frame) rendered results.render()[0] rgb cv2.cvtColor(rendered, cv2.COLOR_BGR2RGB) h, w, ch rgb.shape bytes_per_line ch * w qimg QImage(rgb.data, w, h, bytes_per_line, QImage.Format_RGB888) self.sigFrame.emit(qimg) cap.release() class MainWindow(QWidget): def __init__(self, model): super().__init__() self.model model self.thread None def start_camera(self): self.thread DetectThread(self.model, source0) self.thread.sigFrame.connect(self.update_frame) self.thread.start() def update_frame(self, qimg): self.video_label.setPixmap(QPixmap.fromImage(qimg))这个设计里还有一个容易被忽视的点results.render()是在子线程里调用的绘制检测框这个操作不阻塞UI线程但依然占用推理线程时间。因此推理线程内部不要加time.sleep来限帧直接用摄像头帧率就够。4.3 提升界面流畅度的几个细节为了让UI更流畅我做了几件事一是限制最大帧率。摄像头读帧是8-30fps但模型推理在CPU机器上可能只能跑到5-10fps。如果摄像头帧率远高于推理速度会导致队列积压延时越来越大观感就是画面越来越卡。我在读帧后加了个简单的丢帧策略上一次推理还没结束时本次读到的帧直接丢弃只拿最新帧推理保证延时始终最低。二是不要在主线程做任何图像操作。哪怕是cv2.cvtColor、QImage转换这类看似很快的操作在大量图像帧面前都会造成主线程阻塞。所有图像处理都在推理线程完成主线程只做一次QPixmap.setPixmap。三是启动新线程前要用thread.quit()和thread.wait()清理旧线程。否则反复点“打开摄像头”会创建多个线程同时读取同一个摄像头导致设备冲突界面卡死。4.4 推理集成加载模型与结果解析模型加载部分我放在程序启动时一次性完成不要在每次检测时重复加载。加载用的是torch.hub方式或者直接调用YOLOv5仓库的detect脚本接口。import torch def load_model(weights_path): model torch.hub.load(ultralytics/yolov5, custom, pathweights_path, force_reloadFalse) model.conf 0.5 # 置信度阈值 model.iou 0.45 # NMS IoU阈值 return modelmodel.conf和model.iou的值对检测效果影响很大。conf设太低会出现大量误检比如把黑色的电源线误检成塑料瓶设太高又会漏检小目标。我的调参经验是垃圾检测场景conf 0.5起步如果误检多往上调到0.6如果漏检多往下降到0.4。results.render()方法返回的是画好框的BGR图像数组格式化输出用results.pandas().xyxy[0]每一行包含xmin ymin xmax ymax confidence class name这些字段可以直接用于统计表格。4.5 打包发布与界面显示优化项目完成后如果需要给别人演示推荐用PyInstaller打包pip install pyinstaller pyinstaller -w -F main.py-w是不显示控制台窗口-F是打包成单文件。但打包后体积较大带PyTorch大约1-2GB而且首次启动慢。如果只是本机演示不用打包直接跑脚本更省事。中文显示是另一个容易忽略的坑。YOLOv5默认标签名是英文UI里显示英文没问题但如果代码里写了中文label要确保系统字体支持。PyQt5在Linux上偶尔中文乱码加上QFont设置指定中文字体例如Microsoft YaHei或Noto Sans CJK SC。5. 训练与部署过程中的常见问题排查5.1 训练阶段典型问题速查表问题现象可能原因解决方案torch.cuda.is_available()为FalsePyTorch版与CUDA不匹配/装了CPU版卸载重装对应CUDA的PyTorch版本训练时Windows下DataLoader卡死多进程加载和数据不匹配训练命令加--workers 0loss不降/降得很慢标注错误、学习率不当、数据量不足抽查标注框确认类别对应降低学习率补充数据训练OOM显存爆掉batch-size过大、img-size过大减小batch-size或改--img 416验证集mAP不涨反降过拟合/数据增强过强提前停止训练增加数据调整增强参数类别混淆严重如纸和塑料膜类别边界模糊数据特征重叠增加边界样本、细分特征明显的类别标签5.2 UI界面卡顿问题排查界面卡顿排查按顺序处理先把推理放到子线程再看摄像头帧率、推理帧率是否匹配最后检查主线程是否还有多余的图像处理操作。一个容易忽视的因素是摄像头分辨率。如果摄像头默认1080P而你的模型输入只要640至少一半的时间浪费在读帧和解码上。把摄像头显示分辨率设成1280x720或640x480流畅度提升非常明显。cap cv2.VideoCapture(0) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)5.3 部署与性能优化如果想把系统部署到Jetson Nano这类低算力边缘设备上YOLOv5s在CPU模式下可能只有1-2fps建议走两条路换YOLOv5n小模型或者导出TensorRT加速。python export.py --weights best.pt --include engine --device 0 --imgsz 640TensorRT转换后Jetson Nano上推理速度能提升3-5倍。但TensorRT转换对YOLOv5版本有要求不同版本之间同一个engine文件可能不通用需要在目标设备上重新转换。如果demo要求不高CPU上直接跑YOLOv5n用低分辨率输入如320x320也能维持一个勉强能用的帧率。5.4 一些提高系统稳定性的经验数据统计部分我是把识别结果实时写入SQLite或CSV文件方便做分类占比统计。这个功能处理逻辑很简单但要注意写文件操作不要放在推理线程主路径上用队列缓冲异步写入避免IO阻塞推理。检测结果加上“可回收/其他”的策略规则也很好用。比如识别到“纸板、玻璃、金属、纸张、塑料瓶”UI上直接给“可回收物”结论识别到“厨余”则提示“厨余垃圾”。这样用户不用自己查表项目的实用价值立刻提升一个档次。摄像头异常处理也要做。摄像头被拔掉时cap.read()会返回False不处理的话程序直接崩溃或陷入死循环。我加了个状态检测连续5帧读不到有效画面就弹出提示并自动停止线程这个细节在答辩演示时很加印象分。6. 最后的实战心得与扩展方向整个系统从数据准备到UI集成我断断续续做了差不多三周。感受最深的两点一是数据集的质量决定模型上限标注阶段偷懒后面训练和调参会花几倍时间填坑二是UI卡顿问题不是靠调参解决的而是靠正确的架构——把推理和显示彻底分离系统立马就“活”过来了。做这类项目建议先把一个最小的闭环跑通下载公开数据集、训练100轮出一个基础模型、UI显示一张图片的检测结果然后再逐步加视频流、统计、部署优化。不要在第一天就想着界面做得多么漂亮功能闭环通了界面美化都是后面的加分项。后续要扩展的话几个方向都很值得试接入压力传感器和舵机实现物理分拣把模型导出为ONNX放到Web端跑做成网页扫码上传识别或者换YOLOv8重新训练对比一下精度和速度的差异。这套骨架搭好了往上加功能就是顺水推舟的事。