基于YOLOv5的茶叶目标检测：从数据集构建到树莓派5部署实战

简介本资源面向计算机视觉入门与进阶学习者提供一套基于YOLOv5的茶叶目标检测完整项目实战方案可用于农业智能化场景下的茶叶识别、计数与品质分拣等任务帮助读者掌握从数据配置到模型训练、推理部署的全流程。压缩包共95个文件约242KB以34个Python脚本和41个YAML配置文件为主前者覆盖训练、验证、检测、导出及工具模块后者用于数据集、超参数与模型结构定义另含Shell脚本、Dockerfile、Markdown说明与Jupyter教程便于快速复现与容器化部署。目前已有513人学习下载。项目目录结构清晰包含models、utils、data、scripts等模块读者可据此理解YOLOv5的工程组织方式并借助配套流程教程完成环境搭建、模型训练与效果验证积累目标检测实战经验。1. 茶叶目标检测从嫩芽到成叶YOLOv5 到底能识别到什么程度茶园里最耗人力的环节不是采摘本身而是分级。同一批鲜叶里混着单芽、一芽一叶、一芽二叶还有老叶和茶梗靠人眼在传送带边上盯着分一天下来眼睛发花标准还会漂移。茶叶目标检测要解决的就是这件事让模型在图像里把不同嫩度的茶叶分别框出来输出类别和位置后续接分拣机构或者只做统计估产。基于 YOLOv5 实现茶叶目标检测算法是这条链路里落地成本最低的一条路——单阶段检测、训练快、部署轻树莓派 5 上跑量化后的模型也能到可用帧率。这篇面向的是想拿茶叶数据跑通一套检测流程的从业者你会看到数据集怎么标、YOLOv5 怎么配、训练崩了怎么救、后处理怎么调以及这套方案在真实茶园里到底值不值得投入。2. 茶叶数据集怎么建标注规范、增强策略与目录结构2.1 类别定义先定死别边标边改茶叶检测翻车最常见的原因不是模型不行是类别定义在中途变了。今天把「一芽一叶」和「一芽二叶」分开标明天觉得太细又合并标完的 3000 张图全废。我的习惯是开工前先写一份类别表把每个类别的判定标准用文字锁死再配 5 张典型图作为锚点。以绿茶鲜叶分级为例常见做法是分四类类别 ID类别名判定标准典型易混情况0single_bud只有芽头无展开叶片芽头带极小叶尖仍算单芽1bud_one_leaf一芽一叶第二叶未展开第二叶刚冒尖算一芽一叶2bud_two_leaf一芽二叶第三叶未展开第二叶已完全展开才算3old_leaf老叶、茶梗、破损叶虫咬叶归此类类别数直接决定检测头输出维度。四类够用别一上来标十几类茶叶嫩度是连续量类别越多边界越模糊标注一致性越差。如果后续要做嫩度回归可以先用这四类做粗分再在分类分支上加回归头。2.2 标注工具选型与 YOLO 格式转换目标检测常用标注工具里LabelImg 和 X-AnyLabeling 是茶叶项目里用得最多的两个。LabelImg 轻量、快捷键顺手适合纯手工标X-AnyLabeling 带 SAM 辅助标注对茶叶这种密集小目标能省不少时间。不管用哪个导出时选 YOLO 格式得到每张图一个同名 txt每行是class_id x_center y_center width height全部归一化到 0~1。如果拿到的是 VOC 的 XML需要转一道。下面这个脚本处理单目录下的所有 XML输出到 labels 目录import xml.etree.ElementTree as ET import os # 输入输出目录按实际路径改 voc_dir ./annotations out_dir ./labels os.makedirs(out_dir, exist_okTrue) # 类别名到 id 的映射必须和训练时的 data.yaml 一致 classes [single_bud, bud_one_leaf, bud_two_leaf, old_leaf] for xml_file in os.listdir(voc_dir): if not xml_file.endswith(.xml): continue tree ET.parse(os.path.join(voc_dir, xml_file)) root tree.getroot() size root.find(size) w int(size.find(width).text) h int(size.find(height).text) lines [] for obj in root.iter(object): name obj.find(name).text if name not in classes: continue # 未定义类别直接跳过避免训练时报错 cls_id classes.index(name) bbox obj.find(bndbox) x1 float(bbox.find(xmin).text) y1 float(bbox.find(ymin).text) x2 float(bbox.find(xmax).text) y2 float(bbox.find(ymax).text) # 归一化并转为中心点宽高 xc (x1 x2) / 2.0 / w yc (y1 y2) / 2.0 / h bw (x2 - x1) / w bh (y2 - y1) / h lines.append(f{cls_id} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}) out_name os.path.splitext(xml_file)[0] .txt with open(os.path.join(out_dir, out_name), w) as f: f.write(\n.join(lines))逻辑说明脚本先读图片宽高做归一化再把 VOC 的左上右下坐标转成 YOLO 的中心点加宽高。参数上classes列表的顺序就是类别 id必须和后面data.yaml里的names完全一致顺序错了模型学到的类别就是乱的。if name not in classes这行是后悔药标注时手滑写了错别字这里直接跳过而不是让训练崩掉。2.3 数据增强茶叶场景别乱用翻转YOLOv5 自带 mosaic、HSV 增强、随机缩放和平移。茶叶检测里有两个坑一是上下翻转要慎用茶叶在枝头的朝向有语义上下翻转会造出自然界不存在的样本二是 HSV 的饱和度增强幅度别开太大茶园光照下鲜叶的绿色是重要特征饱和度拉过头会让嫩芽和老叶的颜色区分消失。我一般会在hyp.scratch-low.yaml基础上改三个值flipud设 0.0hsv_s从 0.7 降到 0.4mosaic保持 1.0 但在最后 10 个 epoch 关掉。关 mosaic 是为了让模型在训练末期见到真实分布的图像mAP 通常能回涨一两个点。数据量少于 2000 张时mosaic 的收益最明显超过 8000 张后可以把 mosaic 概率降到 0.5避免过度拼接导致小目标被裁切。目录结构按 YOLOv5 的要求摆dataset/ images/ train/ val/ labels/ train/ val/ data.yamldata.yaml里写train、val的绝对或相对路径nc: 4names按顺序列四类。路径里不要有中文和空格这是血泪经验Windows 上训练时路径带空格会让 dataloader 报一些看不懂的错。3. YOLOv5 环境配置与训练参数从 conda 到第一个 baseline3.1 环境配置conda 建环境加依赖锁定YOLOv5 环境配置的坑集中在 PyTorch 和 CUDA 版本匹配上。我的做法是用 conda 建一个干净环境先装 PyTorch 再装其余依赖不要直接pip install -r requirements.txt让它自己拉 torch。conda create -n tea_yolo python3.9 -y conda activate tea_yolo # 按本机 CUDA 版本选这里以 CUDA 11.8 为例 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 再装 YOLOv5 其余依赖 pip install -r requirements.txt参数说明Python 用 3.9 是因为部分标注和可视化库对 3.10 以上支持还不稳。PyTorch 2.0.1 配 CUDA 11.8 是经过大量项目验证的组合别追最新版。装完用python -c import torch; print(torch.cuda.is_available())验证输出 True 才算通。如果输出 False先查驱动版本再查是不是装成了 CPU 版 torch。3.2 训练命令与关键超参数baseline 训练命令如下python train.py \ --img 640 \ --batch 16 \ --epochs 150 \ --data dataset/data.yaml \ --weights yolov5s.pt \ --cfg models/yolov5s.yaml \ --hyp data/hyps/hyp.scratch-low.yaml \ --name tea_baseline参数逐个说--img 640是输入分辨率茶叶目标偏小如果显存够可以上 960小目标召回会明显提升但训练时间大约翻倍。--batch 16在 8G 显存上跑 640 分辨率比较稳爆显存就降到 8。--weights yolov5s.pt用预训练权重茶叶数据量不大时这是必须的从零训基本收敛不到可用精度。--cfg指定模型结构s 版本够用追求精度换 m 或 l但部署到树莓派 5 上建议还是 s。--hyp用 low 档增强配合前面说的自定义修改。训练过程中重点看三个指标metrics/mAP_0.5、metrics/mAP_0.5:0.95和val/box_loss。mAP_0.5 到 0.85 以上算可用0.9 以上算好。如果 box_loss 震荡不降先查标注框有没有越界或者宽高为 0 的脏数据。3.3 训练崩了怎么救三个高频故障第一个故障是 loss 变 NaN。原因通常是学习率太大或者某批数据里有异常框。解决把lr0从 0.01 降到 0.001同时用脚本扫一遍 labels删掉宽或高小于 0.001 的行。第二个故障是 mAP 卡在 0.3 不动。八成是类别不平衡单芽样本远多于老叶。解决在data.yaml里给老叶类加权或者用--balance类的采样策略简单点就直接复制老叶样本做过采样。第三个故障是验证集 mAP 高但实际推理一塌糊涂。这是过拟合加数据泄漏检查 train 和 val 有没有同图不同增强的重复样本。我一般会按拍摄批次划分 train/val而不是随机分这样验证结果才可信。4. 推理部署与后处理置信度阈值、NMS 和树莓派 5 实测4.1 推理脚本与参数调节训练完拿best.pt做推理python detect.py \ --weights runs/train/tea_baseline/weights/best.pt \ --source test_images/ \ --img 640 \ --conf-thres 0.35 \ --iou-thres 0.45 \ --save-txt--conf-thres 0.35是置信度阈值茶叶密集场景下别设太高0.5 会漏掉大量被遮挡的芽头也别太低0.2 以下老叶误检会暴涨。--iou-thres 0.45控制 NMS 的合并力度同一芽头被重复框时调低这个值能压掉冗余框但太低会把相邻的两个芽合并成一个。--save-txt输出检测框坐标方便后续接分拣逻辑。后处理里还有一个容易被忽略的点YOLOv5 默认的 NMS 是类内抑制不同类别之间不互相抑制。茶叶场景里单芽和一芽一叶位置高度重叠如果发现同一位置出了两个不同类别的框需要在后处理里加类间 NMS或者干脆在训练时把这两类的标注边界再拉开。4.2 树莓派 5 部署模型导出与实测帧率树莓派 5 上部署自己训练的 YOLOv5 模型走 ONNX Runtime 比直接跑 PyTorch 快得多。先导出 ONNXpython export.py \ --weights runs/train/tea_baseline/weights/best.pt \ --include onnx \ --img 640 \ --opset 12--opset 12是兼容性最好的算子集版本树莓派上的 onnxruntime 对更高 opset 支持不一定全。导出后在树莓派 5 上用 onnxruntime 加载输入做归一化和通道转换实测 640 分辨率 yolov5s 大约 8~12 FPS降到 416 分辨率能到 20 FPS 以上。如果还要更快可以走 NCNN 或量化到 INT8但茶叶小目标对量化比较敏感INT8 后 mAP 可能掉 3~5 个点要重新评估。部署时的预处理要和训练时完全一致letterbox 填充、BGR 转 RGB、归一化到 0~1。任何一步不一致推理结果都会系统性偏移这种问题排查起来很费时间建议把预处理参数写进配置文件训练和部署共用。5. 避坑与排查茶叶检测项目里最容易翻车的五件事现象一训练 loss 正常下降但验证 mAP 一直是 0。原因data.yaml里的names顺序和标注时的类别 id 对不上或者 val 路径写错导致加载了空标签。解决打印一次 dataloader 的输出确认标签里的 class_id 都在nc范围内val 目录下图片和 txt 数量一致。现象二模型把老叶全检成单芽。原因单芽样本占比过高模型倾向于把所有绿色小目标都判成多数类。解决统计各类别框数量对少样本类做过采样或在 loss 里加类别权重同时检查标注时老叶和单芽的边界是否清晰。现象三同一片茶叶出好几个重叠框。原因NMS 的 iou 阈值设太高或者茶叶目标本身密集导致 NMS 失效。解决把--iou-thres从 0.45 降到 0.3 试如果还不行考虑换 Soft-NMS 或者 DIoU-NMSYOLOv5 的general.py里可以改 NMS 实现。现象四树莓派上推理结果和 PC 上不一致。原因预处理不一致最常见的是 letterbox 的填充颜色或缩放比例不同以及 BGR/RGB 顺序搞反。解决把 PC 上的预处理输出和树莓派上的预处理输出各存一张图对比逐像素查差异。现象五换一批新茶园的数据mAP 直接掉一半。原因域偏移训练数据的光照、背景、茶叶品种和新场景差异大。解决拿新场景的 200~500 张图做微调冻结 backbone 只训检测头学习率设小一点通常两三个 epoch 就能拉回来。这也是目标检测模型微调崩了之后最有效的补救方式。6. 把茶叶检测做扎实的一个技巧用切片推理救小目标茶叶检测里最头疼的是小目标——单芽在 640 分辨率下可能只占 20 个像素YOLOv5 的 P3 特征图 stride 是 820 像素的目标在特征图上只剩 2~3 个格子信息损失严重。除了前面说的提高输入分辨率还有一个技巧是切片推理SAHI 思路把大图切成带重叠的小块分别推理再合并结果。具体做法是把 4000x3000 的原图切成 640x640 的块重叠 128 像素每块单独送模型最后用 NMS 合并所有块的检测框。代价是推理时间随切块数线性增长但小目标召回能提升 10~20 个点。在估产和分级统计场景里这个提升值得。def slice_inference(img, model, slice_size640, overlap128, conf0.35): h, w img.shape[:2] step slice_size - overlap all_boxes [] for y in range(0, h, step): for x in range(0, w, step): # 边界处理保证切片不越界 x2 min(x slice_size, w) y2 min(y slice_size, h) x1 max(0, x2 - slice_size) y1 max(0, y2 - slice_size) patch img[y1:y2, x1:x2] # 推理并把框坐标映射回原图 results model(patch, confconf) for box in results: box.x1 x1; box.y1 y1 box.x2 x1; box.y2 y1 all_boxes.append(box) # 全局 NMS 合并 return nms(all_boxes, iou_thres0.45)参数上slice_size和训练分辨率保持一致overlap取 slice_size 的 20% 左右太小会漏掉跨块目标太大推理时间浪费。这个方案在茶叶这种密集小目标场景里比单纯堆模型容量有效得多。我自己做茶叶检测项目最大的教训是别在模型结构上折腾太久YOLOv5s 加好数据加切片推理比换个花哨的 backbone 实在。数据标注的一致性、预处理的对齐、后处理阈值的调校这三件事做扎实效果自然出来。希望帮到你。本文还有配套的精品资源点击获取