简介针对茶叶嫩芽目标检测与关键点定位任务这份压缩包提供一套两阶段模型的完整工程化方案适合从事农业视觉、智慧茶园、自动化采摘等方向的算法工程师、研究人员和高校学习者参考复现。方案涵盖YOLO系列或Faster RCNN的目标检测主干以及嫩芽关键点定义与回归映射模块可服务于嫩芽质量评估、生长状态监测和采摘决策等场景。压缩包共1766个文件大小约803.27MB以944个jpg现场图像和795个json标注文件为基础数据配套20个py脚本实现数据预处理、模型训练与推理流程pt、pth、torchscript等模型文件提供预训练权重与部署版本另有bat脚本便于一键启动。已有699人下载学习。内容涵盖图像数据、标注文件、训练脚本与模型权重几大模块可支撑从数据预处理、模型训练到权重部署的完整流程适合作为课程设计、论文实验或工程项目的基线方案也可迁移到其他嫩芽、小目标关键点检测任务中。1. 茶叶嫩芽检测为什么必须走“两阶段”先找到芽再找芽尖做农业视觉落地的人多半遇到过这类场景茶园里的嫩芽和成熟叶片颜色接近、芽体细小一丛茶树枝叶互相遮挡清晨或背光时嫩芽边缘几乎融进背景。拿通用目标检测模型直接上你会发现框能出来但框里的“芽尖”到底在哪个像素上模型完全不关心。可茶园自动采摘、嫩芽分级、芽叶比统计这类需求恰恰需要知道芽尖精确位置和姿态——这是单阶段检测模型给不了的。这个标题里的方案核心思路是两阶段串联第一阶段用目标检测把嫩芽区域定位出来是什么、在哪片区域第二阶段在已定位的区域内做关键点检测输出嫩芽的芽尖、叶柄等关键点位置最后通过“目标-关键点-回归映射”把关键点坐标从区域坐标还原到整幅图像坐标。等于先划重点、再做精读而不是让模型一口气干完所有事。这套路线适合做茶园采摘机器人视觉、嫩芽长势监测、茶叶分级设备的人也适合想在目标检测项目里加入关键点输出的开发者。我按这个标题做下来的经验是两阶段不只是网络结构的堆叠数据标注、损失函数、坐标映射任何一个环节出错模型精度都会卡在某个阈值上不去。下面从模型搭建、标注格式、训练参数、避坑经验到验证手段把整条链路拆开讲。2. 两阶段模型结构目标分支、关键点分支与回归映射的联动方式2.1 为什么不用单阶段模型直接把关键点回归出来有同行问过既然关键点就是几个坐标值能不能让检测头多输出几个通道把芽尖坐标一起回归了我在早期实验里试过这种“一步到位”的做法效果不理想。原因在于嫩芽目标普遍只有三四十个像素在整张茶园图像里占比极小单阶段模型要同时优化“目标在哪”和“关键点在哪”两个任务特征冲突很严重——分类和回归分支会把梯度集中到高频特征上小目标的芽尖位置被平均掉。更麻烦的是关键点坐标是绝对坐标的话模型得为每个目标都输出一组高精度数值训练收敛慢推理时还容易出现关键点漂移。两阶段模型的好处是任务解耦第一阶段只负责“框出嫩芽”第二阶段只在框内做局部关键点回归特征已经对齐过一次任务各自单纯。实际工程里两阶段方案在嫩芽这类小目标上的关键点误差通常比单阶段低30%到50%这是结构本身带来的收益不是调参能弥补的。2.2 头部结构检测头加关键点头怎么把特征“抠”出来常见的做法是共享backbone可选择ResNet50或CSPDarknetbackbone输出特征图后分两条支路。目标检测支路负责框回归和置信度分类关键点支路接收检测支路输出的候选区域Region of Interest在候选区域对应的特征图上做关键点热图回归或坐标回归。我习惯把关键点支路设计成三部分——ROI特征提取、关键点头、坐标映射层。ROI特征提取用PSROIAlign或ROIAlign把不同尺寸的候选区域统一成固定尺寸特征图比如14×14或7×7保证后续全连接层能接收定长输入。关键点头可以是简单的卷积加全连接输出形式有两种直接回归归一化坐标或者输出2D热图再取峰值。下面给一个简化的PyTorch风格结构定义方便理解两个分支的连接关系import torch import torch.nn as nn import torchvision.ops as ops class TeaBudTwoStage(nn.Module): def __init__(self, backbone, num_classes1, num_keypoints2): super().__init__() self.backbone backbone # 比如 ResNet50输出 C4 特征图 self.detection_head nn.Sequential( nn.Conv2d(1024, 256, 3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(256, num_classes * 4 1, 1) # 每个anchor: 4个框坐标 1个目标置信度 ) self.roi_align ops.RoIAlign(output_size(7, 7), spatial_scale1/16, sampling_ratio2) self.keypoint_head nn.Sequential( nn.Linear(1024 * 7 * 7, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.2), nn.Linear(512, num_keypoints * 2) # 每个关键点输出 (dx, dy)相对ROI左上角的偏移 ) def forward(self, images, rois): feat self.backbone(images) # 输入 [B,3,H,W] - 特征 [B,1024,H/16,W/16] det_out self.detection_head(feat) # 检测头输出shape 与 feat 一致 roi_feat self.roi_align(feat, rois) # 按 RoI 裁剪并 resize 到 7x7 b roi_feat.size(0) roi_feat_flat roi_feat.view(b, -1) kpt_offset self.keypoint_head(roi_feat_flat) # 输出相对坐标偏移 return det_out, kpt_offset这段代码里det_out的最后一维是“每个anchor的4个坐标偏移加1个置信度”实际训练时可换成Faster R-CNN或Cascade R-CNN的标准输出头道理一样。rois来自检测分支的候选框格式是[x1, y1, x2, y2]注意RoIAlign要求的是原图坐标系下的绝对坐标而不是归一化坐标。关键参数说明spatial_scale1/16backbone下采样倍数ResNet50的C4输出是原图的1/16如果你的backbone是下采样32倍这里要改成1/32否则特征对齐位置会整体偏移。RoIAlign output_size7统一ROI特征的空间尺寸7×7是常见选择尺寸越大保留的空间细节越多但计算量增大嫩芽目标小我一般用10×10或14×14关键点定位误差更低。关键点头最后一层输出num_keypoints * 2每个关键点两个数表示相对ROI左上角的偏移量。输出后用Sigmoid限制在[0,1]区间再乘以ROI宽高即可得到原图坐标。2.3 目标-关键点-回归映射坐标还原的三个公式这是标题里的核心概念也是很多实现翻车最多的地方。关键点头在ROI坐标系下回归出相对偏移最终必须映射回原图生成绝对坐标才能交给下游的机械臂或统计模块。映射公式如下原图关键点 x roi_x1 offset_x * roi_width 原图关键点 y roi_y1 offset_y * roi_height其中roi_x1是检测框左上角在原图上的x坐标roi_width是检测框宽度offset_x是网络输出的归一化横向偏移0到1之间。如果关键点头输出的是热图而不是偏移映射方式变为先在ROI特征热图上找响应值最大的像素位置再乘以ROI在特征图上的缩放比例最后加上ROI在原图的左上角坐标。两种方式本质相同都是“局部坐标 全局偏移”的组合。这个映射在训练时和推理时各要做一次。训练时把标注的关键点绝对坐标转换成ROI相对坐标作为监督目标推理时反过来把网络输出的相对坐标还原成绝对坐标。务必保证两个过程互逆否则会出“训练时效果好、推理时关键点位置偏半个身位”的诡异问题。3. 数据集与标注格式关键点标注做不好两阶段模型直接失去意义3.1 嫩芽关键点标注定义哪几个点最实用做目标检测时标注只需要一个矩形框但加上关键点后必须先定义清楚“关键点是什么”。我做茶园项目时按采摘需求定义了两类关键点芽尖bud_tip嫩芽最顶端的像素位置是采摘机械臂的切割参考点。芽柄基部bud_base嫩芽与枝干连接处的像素位置用于判断嫩芽朝向和长度也用于计算芽叶比。如果做的是茶树表型分析可以再加“第一片展开叶叶柄点”“第二片展开叶叶柄点”但采摘场景下两个点足够点多反而增加标注不一致的风险。关键点数量建议控制在2到4个每个点都要有明确的像素级定义标注规范里最好配两张示例图正确标注和错误标注。3.2 数据格式从JSON到训练样本的转换脚本常见的关键点数据集格式是MS COCO的扩展格式检测框用[x, y, width, height]表示关键点用[x, y, visibility]数组表示其中visibility取0、1、2分别表示“未标注”“遮挡但存在”“可见”。因为嫩芽经常被叶片遮挡芽基部可能看不到visibility字段一定要利用起来训练时不可见的关键点不参与损失计算。我一般把数据整理成自定义JSON结构如下{ image: tea_bud_001.jpg, width: 1920, height: 1080, annotations: [ { bbox: [623.4, 213.8, 54.2, 41.6], keypoints: [ [640.7, 222.1, 2], [656.2, 241.8, 1] ] } ] }这一段里bbox是检测框左上角坐标和宽高两个关键点分别为芽尖和芽基。visibility1的芽基被遮挡训练时该点的loss置零。下面给一个从标注JSON转到训练输入的脚本片段核心是坐标归一化和映射import json import numpy as np def load_annotation(json_path): with open(json_path, r, encodingutf-8) as f: ann json.load(f) rois, kpt_abs, kpt_rel, kpt_vis [], [], [], [] for item in ann[annotations]: bx, by, bw, bh item[bbox] rois.append([bx, by, bx bw, by bh]) # 转成 [x1,y1,x2,y2] kps np.array(item[keypoints], dtypenp.float32) kpt_abs.append(kps[:, :2]) vis kps[:, 2] # 关键把绝对坐标映射成ROI内的相对偏移 rel (kps[:, :2] - [bx, by]) / [bw, bh] kpt_rel.append(rel) kpt_vis.append(vis) return rois, kpt_abs, kpt_rel, kpt_vis注意脚本里转换的次序先取出绝对坐标和可见性再除以框的宽高得到相对偏移。这里最容易出错的是分母用了图片宽高而不是框宽高——如果按图片宽高归一化网络会同时学习“关键点在图片里的绝对位置”和“关键点在框里的相对位置”两者冲突训练震荡。3.3 数据增强和拼图要注意的事嫩芽数据集通常不大几百张到一两千张图是常态没有预训练权重加持的话两阶段模型学不出来。我一般加载在COCO或ImageNet上预训练的backbone权重然后做在线增强随机翻转注意关键点左右对调、随机旋转10度以内、HSV抖动、高斯模糊模拟阴天。但有一条铁律进行仿射变换或裁剪时检测框和关键点坐标必须同步变换而且翻转时关键点顺序必须重排。比如芽尖和芽基在左右翻转后两个点的左右关系互换如果不对关键点索引做映射训练loss会一直不降复现时最容易卡在这一步。数据量少时还可以复制粘贴小目标从不同图上把嫩芽区域抠出来贴到训练图上做遮挡模拟前提是目标-关键点回归映射保持一致。4. 训练配置与损失函数目标分支和关键点分支的博弈4.1 两个损失怎么配比直接相加会出问题两阶段模型有两个监督信号目标检测损失和关键点损失。常见做法是让关键点损失只作用于被检测出来的ROI但训练初期检测分支本身还不准ROI位置有偏移关键点分支等于在“被污染”的区域上学习容易不收敛。我的做法是分阶段训练。第一阶段冻结backbone的前几层先单独训练检测分支直到mAP达到基础水平第二阶段解冻全部参数加上关键点分支一起联合训练。损失公式如下total_loss w1 * det_cls_loss w2 * det_reg_loss w3 * kpt_losskpt_loss不是对所有ROI都计算只对与真实框IoU大于0.5的ROI计算。w1、w2、w3建议初始值分别设为1.0、1.0和0.5然后观察收敛曲线调整如果关键点loss降不下去说明检测框质量差优先回头修检测分支而不是调loss权重。4.2 关键点损失函数选型SmoothL1还是Wing Loss关键点坐标回归是连续值回归最常见的是SmoothL1 Loss它对离群点不敏感收敛平稳。但嫩芽关键点标注本身有像素级噪声不同标注员可能差两三个像素直接用SmoothL1对小偏差不敏感可能导致模型在细微坐标上“无所谓”。Wing Loss在这里效果更好它对小误差范围的梯度更大专门用于人脸关键点这类高精度任务对嫩芽这种毫米级定位场景也适用。Wing Loss的公式略复杂但PyTorch里手写也就十几行import torch import torch.nn as nn import math class WingLoss(nn.Module): def __init__(self, w10.0, epsilon2.0): super().__init__() self.w w self.epsilon epsilon self.C w - w * math.log(1 w / epsilon) def forward(self, pred, target, visible): diff pred - target diff diff * visible.unsqueeze(-1) # 不可见的关键点不参与 abs_diff diff.abs() loss torch.where( abs_diff self.w, self.w * torch.log(1 abs_diff / self.epsilon), abs_diff - self.C ) return loss.mean()visible是形状为[num_rois, num_keypoints]的掩码值为0或1对应标注里的visibility字段用unsqueeze扩到与坐标维度一致后乘在差值上不可见点位梯度为0。参数说明w控制线性区宽度epsilon控制曲线曲率。对茶叶嫩芽这种像素精度要求高、标注噪声也高的场景我推荐w10、epsilon2如果是遥感或大目标场景w5更合适因为坐标噪声相对目标尺寸没那么多。4.3 训练脚本关键参数清单训练配置直接决定模型能不能收敛。下面是我在单卡NVIDIA RTX 3090上验证过的一组合适参数供参考参数推荐值说明输入分辨率1333×800原图16:9缩放保小目标细节batch size4两阶段模型显存占用高batch过大易OOM基础学习率0.001加载预训练权重时从1e-4起步更稳学习率衰减每10个epoch乘以0.1训练到20-25轮会明显下降关键点损失权重0.5起步联合训练时不要一开始给太高ROI采样数128正负样本比1:3保证关键点分支看到足够正样本NMS阈值0.5两阶段模型内部处理候选框用最大epoch50嫩芽数据集1000张图50轮可收敛训练时建议把三个loss分别打印出来记录到日志不要只盯total loss。如果kpt_loss在联合训练初期反而上升说明检测分支还在调整框位置关键点分支被迫跟着学这种抖动在前5-10个epoch内正常超过15轮还在抖检查ROI采样逻辑和数据标注一致性。4.4 目标-关键点-回归映射在训练脚本里的实现槽位训练代码里对应关系是这样组织的detection_head输出候选框进NMS后得到roisrois和真实关键点一起交给数据加载器训练时直接用真实框做RoIAlign称为“正样本ROI训练”推理阶段则换成检测预测的框做RoIAlign。两者必须用同一套映射公式否则模型在训练和推断时看到的输入分布不一致典型表现是验证集loss很低但实际推理关键点位置发飘。5. 避坑与常见问题排查嫩芽目标小坑比想象中多5.1 现象训练损失下降但关键点整体偏移半厘米原因回归映射公式写错了最常见的是RoIAlign的spatial_scale和backbone实际下采样倍数不一致。比如backbone输出stride32但代码里写1/16ROI特征被拉伸了两倍关键点上采样回原图时坐标全体偏移。解决打印backbone对单张图的输出尺寸算出实际stride再做映射或者干脆用end2end方式从backbone向前传一次确定输出shape不要在配置里手写hardcode。5.2 现象两个重叠嫩芽检测框一个大一个小关键点经常串位原因嫩芽彼此紧邻时NMS会保留两个有重叠的框第二个阶段同时对两个ROI提取特征如果RoIAlign的采样点太少相邻嫩芽的特征混在一起关键点头区分不开。解决把RoIAlign的sampling_ratio从2提高到4增加采样点密度同时把NMS阈值从0.5降至0.4减少重叠框数量。如果还不行在关键点头前加一个小型自注意力模块让每个ROI内部的特征关注到该ROI对应的嫩芽本体。5.3 现象训练集mAP不错但傍晚和逆光条件下关键点全部漂移到叶片边缘原因训练数据里缺少弱光场景。嫩芽在逆光下是半透明的嫩黄色边缘特征弱模型只能靠形状先验猜芽尖位置一旦光照分布偏移先验失效。解决在数据增强里增加亮度扰动和局部阴影模拟把HSV的V通道随机乘0.6到1.4再加一个随机的径向渐变模拟侧光和背光另外单独收集一个时段的弱光子集做微调。这类问题没有模型结构上的捷径只能把数据分布补齐。5.4 现象visible1可见但遮挡的关键点loss在训练时反而干扰正常点原因visible1的点label位置是人工估计的标注员在遮挡情况下画的位置偏移比其他点大SmoothL1或Wing Loss给这些模糊点分配正常梯度模型被迫拟合错误监督。解决把visible1和visible2分开处理。visible1的点loss权重降为0.3visible2正常参与。实现上在构建visible掩码时额外乘一个权重因子即可代价极小但收敛稳定性明显提升。5.5 现象训练时OOM显存爆掉原因两阶段模型在训练阶段要为每个ROI单独做RoIAlign和关键点计算batch里候选框多时显存开销暴涨尤其是用了14×14以上output_size。解决减小batch size到2同时限制每张图最多保留256个候选ROI如果还爆把关键点头的全连接层改成1×1卷积加全局平均池化可省掉约三分之一显存。不要在模型结构已经写死之后硬调batch先剪ROI数量再降输出尺寸。6. 验证与进阶用OKS评估关键点让模型真正落到茶园6.1 检测与关键点分别算指标别混在一起两阶段模型要分开评估。检测部分看AP0.5和AP0.75判断定位精度是否够用关键点部分用OKSObject Keypoint Similarity评估公式里需要设定每个关键点的标准差sigma芽尖和芽基的sigma我设为20像素对应茶园图像上约0.5厘米的允许误差。OKS阈值为0.75时关键点精度才算“可被机械臂使用”。评估脚本可以直接用COCO API的cocoEval把标注转成COCO格式后调用即可。另外务必统计推理耗时两阶段模型检测加关键点回归在GPU上一般30-50ms但到嵌入式端会到100ms以上采摘场景通常要求单帧80ms以内提前做好量化或换轻量backbone的预算。6.2 进阶技巧把关键点收益反馈到检测分支联合训练稳定后可以做一次可选的refine用关键点预测结果修正检测框。具体做法是计算芽尖和芽基的连线外扩一定比例得到更贴合嫩芽姿态的新框把新框重新送进ROIAlign后再预测一次关键点迭代一次。这个操作能提升关键点整体OKS约两个点代价是推理时多一次前向。我在茶园实地跑这套方案时发现最影响落地效果的不是mAP而是关键点的时域稳定性——同一片叶片在连续帧里抖动几像素机械臂就会来回颤。一个实用的技巧是输出前用卡尔曼滤波平滑芽尖坐标以芽尖历史位置做预测、当前帧检测结果做观测噪声能降一半以上。这不属于模型结构问题但比换模型更直接有效。6.3 给关键点预测画热力图先于标注检查错误提高验证效率有个好手段把关键点头的中间层热力图可视化出来叠在原图上。如果热力图高响应区域不是芽尖而是叶缘或茶枝说明模型学到了错误上下文这时不要盲目调参先检查标注图形状是否一致再决定是增数据还是约束映射分支。两阶段模型的结构收益都在局部特征里热力图是唯一能直观看到模型内部选择的窗口。做这个方向最深的教训是茶叶嫩芽的数据永远不会干净遮挡、逆光、露珠反光全都是常态。两阶段模型解决的是“目标本身太小、关键点需要高精度”的结构问题但数据质量决定精度的天花板。先花时间统一标注规范、补齐弱光场景样本再谈改网络结构不会走冤枉路。希望这套目标-关键点-回归映射的实践记录能帮你在茶园场景里少踩几个坑。本文还有配套的精品资源点击获取