茶叶嫩芽检测两阶段模型：目标-关键点-回归映射全解析

简介面向茶叶生产智能化和计算机视觉研究者压缩包内提供了一套完整的茶叶嫩芽目标与关键点两阶段检测方案。方案将目标检测与关键点定位相结合适用于嫩芽质量评估、自动采摘定位等场景目标检测阶段可选用YOLO系列或Faster RCNN训练中注重不同光照、角度和生长阶段的样本并用精确率、召回率、F1值等指标评估。关键点阶段则定义嫩芽关键点通过回归映射实现精确定位。包内含1766个文件包括944张茶叶嫩芽图像、795个JSON标注文件、20个Python脚本、3个PT与2个PTH模型权重、1个TorchScript模型及1个训练启动脚本整体约803MB兼顾训练与推理需求。已有699人学习下载。通过源码、预训练权重和标注数据读者可以复现完整流程理解“目标-关键点-回归映射”的建模思路并进行迁移训练或二次开发。1. 茶叶嫩芽检测为什么要两阶段从目标框到关键点的回归映射茶叶嫩芽的目标检测其实不难难的是把“芽尖”这个具体的采摘点一起定位出来。机械臂摘茶、无人机数芽、表型测量仪测芽长这些场景不只需要“这里有一片嫩芽”这个框还需要知道芽尖落在像素的哪个坐标上。换句话说先框出嫩芽再在框内找关键点最后把这两者用回归映射钉在一起——这正是这份“茶叶嫩芽目标及其关键点检测两阶段模型目标-关键点-回归映射.zip”资源包解决的完整链路。适合正在做农业视觉项目、或者想在自己数据集上复现两阶段检测加关键点回归的工程师。这篇笔记会把模型结构、数据格式、训练参数和踩坑点全部拆开讲让你拿到就能跑、跑完敢改。2. 两阶段模型架构拆解检测分支与关键点回归分支的分工两阶段模型的思路很直白第一个阶段把茶叶嫩芽当作普通目标检测来做输出目标框第二个阶段只针对第一阶段提出的候选区域做精细的关键点回归输出芽尖在这个目标框内部的相对位置。2.1 第一阶段目标检测器的选型与输出定义这个阶段的核心是“找得准”。茶叶嫩芽在自然光照下和老叶、枝干的边界往往不清晰如果第一阶段就把框漏了第二阶段再精细也白搭。资源包里第一阶段用的是类 YOLO 风格的轻量检测头骨干网络可以选择 MobileNetV3 或者 ResNet18 这一档。这类骨干在嵌入式设备上跑得动同时特征提取能力对茶叶嫩芽这种纹理不算复杂的目标足够用。第一阶段输出的是一个五元组目标的类别置信度、目标框中心点坐标 (cx, cy)、目标框宽高 (w, h)。注意这里目标框中心点坐标用的是相对图像尺寸的归一化值不是绝对像素值。归一化的好处是不同分辨率输入下都不用改回归目标的数值范围模型训练更稳定。class Stage1Detector(nn.Module): def __init__(self, backbone_out_channels256): super().__init__() # 骨干网络输出的特征图经过一个 3x3 卷积后分叉 self.shared_conv nn.Conv2d(backbone_out_channels, 128, 3, padding1) # 类别分支区分背景和嫩芽 self.cls_conv nn.Conv2d(128, 1, 1) # 回归分支输出 cx, cy, w, h全部是归一化数值 self.box_conv nn.Conv2d(128, 4, 1) def forward(self, features): x F.relu(self.shared_conv(features)) cls_logits self.cls_conv(x) box_pred self.box_conv(x) # [B, 4, H, W] return cls_logits, box_pred逻辑说明这里把检测头拆成三条独立卷积分支类别分支和回归分支共享前面一个 3x3 卷积提取的中间特征。共享卷积能减少参数量也强制两个分支从同一组特征出发避免训练时各自漂移。box_pred 输出的四个通道按顺序对应 cx, cy, w, h每个位置的预测值都经过 sigmoid 归一化到 0 到 1。参数说明backbone_out_channels 取决于你选的骨干网络最后输出的通道数ResNet18 是 512MobileNetV3 是 576资源包代码里留了适配层你只需要改这一个入口参数。cls_conv 输出 1 个通道是因为二分类嫩芽 / 背景如果你要同时检测嫩芽和芽苞两个类别改成 2。2.2 第二阶段关键点回归头与归一化坐标设计第二阶段的输入不是完整的特征图而是第一阶段目标框在原始图像上的对应区域。从原图上裁剪出目标区域后缩放到固定尺寸再送入一个小型回归网络。这里的关键设计是关键点输出的是目标框左上角到芽尖的相对偏移量用框宽高做归一化。这样模型不需要记住绝对值只需要学会在“这个框内的这个位置”上找芽尖。class Stage2KeypointHead(nn.Module): def __init__(self, input_size64): super().__init__() # 把裁剪区域从 [B, 3, input_size, input_size] 压成特征向量 self.feature_extractor nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2) ) # 经过两层全连接后输出 dx, dy self.fc nn.Sequential( nn.Linear(64 * (input_size // 4) * (input_size // 4), 128), nn.ReLU(inplaceTrue), nn.Linear(128, 2) # 关键点偏移量 ) def forward(self, cropped_patches): x self.feature_extractor(cropped_patches) x x.view(x.size(0), -1) offset torch.sigmoid(self.fc(x)) return offset # dx, dy 都在 [0, 1] 范围逻辑说明裁剪区域被统一缩放为 64x64经过两次池化后空间尺寸变成 16x16展平成向量后经全连接层输出两个值。sigmoid 把输出限制在 0 到 1 之间表示芽尖相对目标框左上角的横向和纵向比例。比如输出 (0.5, 0.5) 就意味着芽尖在目标框的正中心。参数说明input_size 是裁剪区域的统一边长64 是初始数值如果你的数据集里嫩芽目标偏大或偏小可以按目标框的平均尺寸来调。比如目标框平均边长是 120 像素建议把 input_size 调到 96 或 128避免缩小后丢失纹理细节。2.3 目标框与关键点的映射从相对偏移还原像素坐标两阶段模型最终的预测结果不是第二阶段输出的 dx 和 dy而是把它们映射回原始图像坐标系。映射公式是关键点像素坐标等于目标框左上角坐标加上归一化偏移量与框宽高的乘积。这个过程在推理阶段完成训练阶段拿到的是归一化目标所以损失函数可以直接计算。def map_keypoint_to_image(box_pred, kpt_offset): box_pred: [B, 4] 格式为 [cx, cy, w, h]全部归一化 kpt_offset: [B, 2] 格式为 [dx, dy]相对于目标框左上角 # 先把中心点格式转成左上角格式 x1 (box_pred[:, 0] - box_pred[:, 2] / 2).unsqueeze(1) y1 (box_pred[:, 1] - box_pred[:, 3] / 2).unsqueeze(1) w box_pred[:, 2].unsqueeze(1) h box_pred[:, 3].unsqueeze(1) # 还原关键点像素坐标仍为图像归一化 kpt_x x1 kpt_offset[:, 0:1] * w kpt_y y1 kpt_offset[:, 1:2] * h return torch.cat([kpt_x, kpt_y], dim1)逻辑说明这个映射函数是“目标-关键点-回归映射”名字里“映射”二字的实体。目标框提供空间范围关键点回归分支提供相对位置两者经过一次线性计算得到最终的芽尖坐标。这个映射写成独立函数的好处是无论训练还是推理只要走同一个函数就不会出现两阶段坐标体系不一致的问题。参数说明box_pred 里的 w 和 h 是归一化宽高最大为 1。所以还原出的 kpt_x 和 kpt_y 也是归一化坐标后续可视化时需要再乘以图像宽度和高度。3. 数据标注格式与预处理让回归映射有可靠的真值模型架构只是骨架数据真值的构造方式直接决定回归映射能不能收敛。这个环节要解决三个问题手工标注结果如何组织成模型能读的格式、关键点真值如何从像素坐标转换成目标框内偏移量、训练时如何做数据增强而不破坏关键点与目标框的对应关系。3.1 JSON 标注格式设计与坐标归一化资源包里的标注文件使用 JSON 组织每张图像一条记录包含图像路径、尺寸、目标框列表和关键点列表。目标框记录为左上角和右下角两个点的像素坐标关键点记录为芽尖的像素坐标。训练前需要把像素坐标归一化到 0 到 1避免不同分辨率图像给回归分支带来尺度不一致的问题。import json def parse_annotation(json_path): with open(json_path, r, encodingutf-8) as f: data json.load(f) samples [] for img_info in data[images]: width img_info[width] height img_info[height] for ann in img_info[annotations]: # 目标框为 [x1, y1, x2, y2]像素坐标 x1, y1, x2, y2 ann[bbox] # 关键点为 (kx, ky)像素坐标 kx, ky ann[keypoint] # 归一化目标框 cx ((x1 x2) / 2) / width cy ((y1 y2) / 2) / height w (x2 - x1) / width h (y2 - y1) / height # 关键点转成相对于目标框左上角的偏移 dx (kx - x1) / (x2 - x1 1e-6) dy (ky - y1) / (y2 - y1 1e-6) samples.append({ img_path: img_info[file_name], box: [cx, cy, w, h], kpt: [dx, dy] }) return samples逻辑说明代码的核心是把像素坐标往归一化空间转换。cx 和 cy 是目标框中心点在图像上的相对位置dx 和 dy 是关键点相对目标框左上角的偏移比例。注意分母加了 1e-6 防止目标框宽高为零时除零报错。参数说明bbox 的格式必须是 [x1, y1, x2, y2]如果你的检测工具导出的是 [x, y, width, height] 格式需要先转换成 x2 x width 再传入。关键点只取一个芽尖如果你的任务还想加一个叶柄点把 keypoint 改成二维数组后续模型输出通道数也要从 2 改成 4。3.2 数据增强策略翻转和旋转时关键点真值同步变换茶叶嫩芽在田间的姿态比较随机水平翻转和 90 度旋转是性价比最高的两种增强。但这两类操作都会改变目标框和关键点的坐标关系增强时不能只对图像做变换必须同步更新标注真值。资源包代码里直接用仿射变换矩阵把所有标注一起变换避免手工逐项更新时漏掉某个目标。import cv2 import numpy as np def augmented_sample(img, box_pixel, kpt_pixel, flip_hTrue): h, w img.shape[:2] if flip_h and np.random.rand() 0.5: img cv2.flip(img, 1) x1, y1, x2, y2 box_pixel kx, ky kpt_pixel # 新目标框水平翻转后 x 坐标镜像 new_x1 w - x2 new_x2 w - x1 # 新关键点 x 坐标同样镜像y 不变 new_kx w - kx box_pixel [new_x1, y1, new_x2, y2] kpt_pixel [new_kx, ky] return img, box_pixel, kpt_pixel逻辑说明水平翻转的镜像规则是每个像素的 x 坐标变成 宽度减去原来的 x 坐标。目标框的左上角和右下角都要做同样的变换交换后 new_x1 一定小于 new_x2。关键点的 y 坐标不变因为水平翻转不改变纵向位置。参数说明flip_h 设为 True 代表启用水平翻转增强。实际训练时建议同时做随机亮度和对比度调整模拟阴天和晴天不同光照但这两类颜色增强不需要修改标注直接套在图像上就行。3.3 训练数据目录结构与验证集划分资源包的数据目录按 train / val / test 三个子目录组织每个子目录下是图像文件对应标注文件统一放在 annotations 文件夹下。划分数据时要保证同一株茶树的图像不会同时出现在训练集和验证集里否则验证分数虚高。茶树的嫩芽和母株在图像上的位置相对固定简单随机划分容易让模型记住位置特征而不是嫩芽本身的纹理。import os import random def split_dataset(all_files, train_ratio0.7, val_ratio0.15): random.seed(42) # 固定随机种子保证可复现 random.shuffle(all_files) total len(all_files) train_end int(total * train_ratio) val_end int(total * (train_ratio val_ratio)) train_files all_files[:train_end] val_files all_files[train_end:val_end] test_files all_files[val_end:] return train_files, val_files, test_files逻辑说明固定随机种子是数据划分的关键操作。不固定的话每次跑训练脚本数据分布都不一样模型在验证集上的分数无法横向对比。如果你希望按茶树个体划分先按图像文件名里的茶树编号分组再对组做随机划分。参数说明train_ratio 和 val_ratio 分别控制训练集和验证集比例剩余部分自动做测试集。嫩芽数据量少于 1000 张时建议把 test_ratio 压低到 15%把更多数据留给训练测试时用手动拍摄的新场景验证泛化性。4. 训练配置与参数调优让回归映射真正收敛两阶段模型的难点在训练策略不是单纯堆训练轮数。第一阶段和第二阶段的任务难度不同收敛速度也不一样。如果直接用同一学习率从头训练到底第二阶段关键点回归分支很容易被第一阶段的目标框损失淹没。资源包采用分阶段训练策略外加独立的损失权值控制。4.1 分阶段训练策略与学习率安排第一阶段先冻结关键点回归分支只训练目标检测部分。目标检测收敛后再解冻第二阶段把第一阶段的学习率调低一个数量级专心优化关键点回归。这样做的原因是第一阶段不稳定时它输出的目标框位置和大小一直在抖第二阶段拿到的是不断变化的输入区域永远学不稳。def configure_optimizer(model, stage): if stage 1: # 只更新第一阶段检测头的参数 params list(model.stage1.parameters()) lr 1e-3 elif stage 2: # 第一阶段用很小的学习率微调第二阶段用正常学习率 params [ {params: model.stage1.parameters(), lr: 1e-5}, {params: model.stage2.parameters(), lr: 1e-3} ] lr None # 每个参数组单独指定 optimizer torch.optim.AdamW(params, lrlr if lr is not None else 1e-3) return optimizer逻辑说明阶段 1 的优化器只包含第一阶段检测头的参数骨干网络如果也是随机初始化可以一并放进去。进入阶段 2 后第一阶段用 1e-5 的学习率做微调只让目标框稍微修正不至于剧烈移动第二阶段用 1e-3 正常学习率快速收敛。这里用 AdamW 而非 SGD因为关键点回归分支的学习信号偏弱AdamW 的自适应学习率能更快找到合适步长。参数说明1e-3 作为阶段 1 的初始学习率适用于大多数图像尺寸为 640 的训练任务。如果你的输入分辨率是 1280批量大小不变的情况下梯度噪声变小学习率可以适当提高到 2e-3。4.2 损失函数设计分类、定位、关键点的权重配比训练时总损失由三部分构成目标检测的分类损失、目标框的回归损失、关键点回归损失。前两个损失由第一阶段计算关键点损失由第二阶段计算。默认权重配置是分类损失权重 1.0、目标框回归损失权重 1.0、关键点回归损失权重 2.0。关键点权重要高一倍因为第二阶段的任务更难且在训练初期梯度信号更弱。import torch.nn.functional as F def compute_total_loss(cls_logits, box_pred, kpt_pred, targets): # 第一阶段类别损失用二分类交叉熵 loss_cls F.binary_cross_entropy_with_logits( cls_logits, targets[cls] ) # 第一阶段目标框回归损失用 smooth L1 loss_box F.smooth_l1_loss(box_pred, targets[box]) # 第二阶段关键点回归损失用 smooth L1权重更高 loss_kpt F.smooth_l1_loss(kpt_pred, targets[kpt]) total_loss loss_cls loss_box 2.0 * loss_kpt return total_loss, { loss_cls: loss_cls.item(), loss_box: loss_box.item(), loss_kpt: loss_kpt.item() }逻辑说明smooth L1 损失对离群点的梯度压制比 MSE 好当预测值和真值差得很远时不会产生爆炸梯度。这在关键点回归里很关键——训练初期模型经常输出离谱的偏移值用 MSE 会让梯度大得直接冲坏整个网络。分类损失用的是带 logits 的 BCE内部已经做了 sigmoid不需要在模型输出后再单独做一次。参数说明2.0 这个权重不是固定的如果训练过程中 loss_kpt 一直降不下去而 loss_box 已经收敛可以把关键点权重提到 3.0 或 4.0。反之如果关键点抖动明显说明权重过高导致关键点分支过拟合应该降回 1.5。没有万能参数每次调整后至少观察 20 个 epoch 再下结论。4.3 epoch 数量、批量大小与验证频率资源包默认配置是阶段 1 训练 150 个 epoch阶段 2 训练 200 个 epoch总轮数偏多是考虑到茶叶嫩芽数据集通常不大。批量大小建议设成 16如果 GPU 显存不够就调到 8同时把学习率减半。验证建议每个 epoch 都做一次在验证集上同时计算目标框的 mAP 和关键点的平均像素误差。关键点误差的衡量需要单独说明把验证图像里的关键点预测值和真值都还原成像素坐标计算两者的欧氏距离再求平均。这个指标直接反映机械臂采摘点定位的精度比只看损失值更直观。200 个 epoch 跑完后这个平均像素误差应该在 5 像素以内才算合格。如果远超 5 像素回查第二阶段输入裁剪区域的大小设置大概率是目标区域被缩放后芽尖纹理丢失过于严重。5. 常见问题与避坑回归映射训练翻车记录两阶段模型跑起来不算难但要稳定收敛需要处理几个经典问题。以下是我拆这个资源包时实际踩过的坑每条按现象、原因、解决的思路记录。5.1 关键点输出值总在 0.5 附近徘徊现象训练前几个 epochloss_kpt 下降得很快但把预测值打印出来发现 dx 和 dy 全部接近 0.5就像是模型每张图都在猜“芽尖在目标框正中间”。原因sigmoid 输出的初始值就是 0.5梯度下降一开始会把所有预测值往训练数据的均值方向拉。如果数据集里大部分芽尖恰好分布在目标框中心附近模型就直接躺平在这个均值上不再学习真实分布。解决给第二阶段回归头添加一个小的偏置初始化扰动把最后一层线性层的 bias 设置成服从均值为 0、标准差为 0.1 的正态分布。同时把训练数据里目标框中心区域 50% 以上的样本做一次随机裁剪偏移让芽尖在目标框内的位置更多样强迫模型学习真实规律而非偷懒猜均值。5.2 第一阶段目标框和第二阶段关键点失配现象训练完成后单看第一阶段输出目标框很准单看第二阶段在真值框上预测的关键点也很准但串联推理时关键点大幅偏离芽尖。原因第二阶段训练用的裁剪区域是标注真值目标框推理时用的是第一阶段预测目标框。预测框和真值框存在偏差哪怕只偏几个像素第二阶段输入区域的分布就跟训练时不一样。解决训练阶段二时对真值目标框加上随机扰动扰动量取目标框宽高的 10% 到 20%让第二阶段在略微偏移的输入上也能给出正确关键点。这相当于给第二阶段做输入区域的域随机化强制它忽略目标框边缘的微小位置误差。资源包代码里这块实现为 jitter_ratio 参数默认设 0.15 效果良好。5.3 回归数值溢出导致 loss 变为 NaN现象训练到第 30 个 epoch 左右loss 突然变成 NaN重新加载保存的模型后再训也很快复现。原因日志里看到 box_pred 的某个通道值在梯度更新后超过了 1下一轮计算 smooth L1 时数值进一步膨胀。通常是因为学习率偏大或者某个样本的标注里存在异常值比如关键点恰好落在目标框外导致标注真值 dx 或 dy 大于 1梯度一直往边界推。解决在 dataloader 里对标注样本做合法校验dx 或者 dy 超出 0 到 1 范围的样本直接剔除并打印警告信息。另外在训练脚本里用 clip_grad_norm_ 限制最大梯度范数为 10这是兜底措施防止任何异常样本直接冲崩模型。两个措施同时加基本能杜绝 NaN 问题。5.4 背景或老叶被误检成嫩芽并参与关键点回归现象验证集上 precision 不低但关键点平均误差非常大。抽查预测结果发现有不少关键点是在误检目标上预测的那些地方压根不是嫩芽。原因第一阶段把老叶或枝干上的高光区域误检成嫩芽第二阶段不管三七二十一对任何输入裁剪区域都会输出一个关键点坐标。关键点分支不知道输入是不是真正的目标它只负责回归相对位置。解决给第二阶段回归头增加一个质量估计输出用目标检测里常用的 IoU 预测方式让关键点分支同时输出“这个裁剪区域包含有效目标的概率”最终判断时把这个概率乘到检测置信度上。概率低于阈值的检测直接丢弃不进入关键点映射环节。这个方法比单纯把阈值调高更有效因为有些真正目标检测置信度其实很低只是关键点质量值会暴露出来。6. 验证与部署可视化脚本与关键点精度检查模型训完第一件事不是测速度而是把预测结果可视化回原图上用眼睛确认芽尖点是不是画在了正确的位置。我一般会在推理脚本里加一个保存可视化结果图的开关每跑完一次预测就从验证集里随机抽 20 张图检查三件事目标框有没有漏检、关键点是否落在芽尖而非芽叶、以及框内是否存在多个嫩芽但只给了一个关键点。import cv2 import torch def visualize_inference(model, img_path, devicecuda): img cv2.imread(img_path) h, w img.shape[:2] img_tensor torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0 img_tensor img_tensor.to(device) with torch.no_grad(): cls_logits, box_pred model.stage1(img_tensor) # 提取置信度大于 0.5 的目标框 scores torch.sigmoid(cls_logits.squeeze(0).squeeze(0)) # 这里简化处理实际要按空间位置聚合 for cx, cy, w_n, h_n in box_pred.squeeze(0).cpu().numpy(): x1 int((cx - w_n / 2) * w) y1 int((cy - h_n / 2) * h) x2 int((cx w_n / 2) * w) y2 int((cy h_n / 2) * h) # 裁剪区域送入第二阶段 crop img_tensor[:, :, y1:y2, x1:x2] crop F.interpolate(crop, size(64, 64), modebilinear) kpt_offset model.stage2(crop) kx int((x1 kpt_offset[0, 0].item() * (x2 - x1))) ky int((y1 kpt_offset[0, 1].item() * (y2 - y1))) # 画框和关键点 cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.circle(img, (kx, ky), 4, (0, 0, 255), -1) cv2.putText(img, f({kx},{ky}), (kx 8, ky - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) cv2.imwrite(pred_ img_path.split(/)[-1], img)逻辑说明代码里第一阶段的输出直接遍历所有空间位置实际工程中要做一次简单的置信度筛选和非极大值抑制否则同一目标会重复画框。关键点坐标还原用的是前面第 2.3 节里定义的映射逻辑先取目标框左上角像素坐标再加归一化关键点偏移量与目标框宽高的乘积。参数说明scores 的阈值 0.5 是初始参考值。如果可视化结果里漏检多阈值往下调收益有限更重要的是看训练集里是不是有大量模糊目标如果误检多优先调整的是第 5.4 节提到的质量估计分支而不是直接调这个绘制阈值。部署到实际采摘或测量场景时要注意推理代码里使用与原图等比例缩放后的输入图输出关键点坐标要乘回原图尺度。模型在 640 分辨率下训练实际摄像头输出 1920x1080 的图直接推理没问题但关键点坐标映射时用的是预测图上的宽高比例不是原图绝对像素值必须在后处理里乘一个 s 原图尺寸 / 模型输入尺寸 的比例系数。这个细节漏掉的话框位置对但关键点在原图上整体偏左上排查成本不低。从那以后我每次调完模型都会强制自己走一遍“可视化回原图 → 目测 20 张 → 再调参”的流程不再只盯着损失曲线和 mAP 数字。这个习惯帮我躲开了好几轮无效调参也让我这个资源包里的可视化脚本持续用到现在每次换数据集都先跑它确认标注没跑偏。希望帮到你少走我走过的弯路。本文还有配套的精品资源点击获取