训练自己的装甲板或能量机关模型

这一节完成一个最小但完整的训练闭环：准备数据、定义关键点、训练 YOLO Pose、检查错误样本并导出 ONNX。为了让流程容易复现，入门实验先只训练一种目标；跑通后再增加颜色、编号和更多场景。

1. 先选择任务形式

只训练目标框

如果目标只是判断画面中有没有装甲板，可以使用普通目标检测。标注成本低，但检测框不能直接作为高质量 PnP 输入。

同时训练目标框与关键点

自瞄更推荐关键点检测：

装甲板：4 个关键点；
能量机关：根据解算定义标注固定的 5 个关键点；
每个数据集中的关键点数量、含义和顺序必须统一。

入门装甲板模型可以只设置一个类别 armor，先验证“能否稳定找到四角点”。颜色和编号可以在后续通过更多类别、额外分类头或单独分类器加入。

2. 建立训练环境

以下命令以 Linux 为例：

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install ultralytics opencv-python
yolo checks

如果需要 NVIDIA GPU，先按照 PyTorch 官方安装页面选择与本机驱动匹配的 CUDA 版本，再安装 Ultralytics。用下面的命令确认训练框架看到 GPU：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU')"

3. 采集数据：先覆盖变化，再增加数量

数据需要覆盖模型实际会遇到的变化：

不同距离、位置、俯仰角和旋转角；
室内、室外、逆光、局部过曝与低曝光；
云台运动造成的方向性模糊；
完整可见、部分遮挡、出画和灯条粘连；
红方、蓝方、灰色和紫色装甲板；
场地灯、反光、屏幕、文字等容易误识别的负样本；
模拟器图像与真实相机图像。

连续视频不要随机按帧切分训练集和验证集。相邻帧几乎相同，随机切分会让验证结果虚高。更合适的方式是按录制场次、地图、光照或机器人划分数据，使验证集和测试集包含训练阶段没有见过的连续片段。

一个可用的初始比例是：训练集 70%、验证集 20%、测试集 10%。比例不是重点，数据来源之间真正独立才是重点。

4. 定义关键点顺序

本教程的装甲板示例统一使用：

0 左下    1 左上    2 右上    3 右下

能量机关也需要先写出类似的语义表，例如中心点、目标叶片特征点和方向点。标注人员必须按语义点选，不能根据点在当前图像中的左右位置临时改变编号。

水平翻转装甲板时，关键点映射为：

[0, 1, 2, 3] → [3, 2, 1, 0]

错误的 flip_idx 会让增强后的样本拥有错误标签，模型可能出现“框很准但角点交叉”的现象。

5. 数据目录与 YAML

目录可以组织为：

datasets/armor_pose/
├── images/
│   ├── train/
│   ├── val/
│   └── test/
└── labels/
    ├── train/
    ├── val/
    └── test/

创建 armor-pose.yaml：

path: /absolute/path/to/datasets/armor_pose
train: images/train
val: images/val
test: images/test

names:
  0: armor

kpt_shape: [4, 3]
flip_idx: [3, 2, 1, 0]

YOLO Pose 的一行标签包含类别、边界框和所有关键点：

class cx cy w h x0 y0 v0 x1 y1 v1 x2 y2 v2 x3 y3 v3

坐标通常归一化到 [0, 1]；v 表示关键点的可见状态。标注完成后，先随机绘制几十张标签图，人工检查框、角点顺序、遮挡点和越界坐标，再开始训练。

6. 训练一个关键点模型

以当前 Ultralytics Pose 接口为例，可以从轻量预训练模型开始：

yolo pose train \
  model=yolo26n-pose.pt \
  data=armor-pose.yaml \
  epochs=150 \
  imgsz=640 \
  batch=-1 \
  device=0 \
  workers=8 \
  seed=42 \
  project=runs/armor_pose \
  name=baseline

n 模型适合先验证数据和标签；数据流程正确后，可以比较 s 或其他规模。模型越大不一定越适合上车，最终需要在 NX 或计划使用的计算平台上测量延迟。

如果显存不足，可以减小 batch 或 imgsz：

yolo pose train model=yolo26n-pose.pt data=armor-pose.yaml epochs=150 imgsz=512 batch=8 device=0

能量机关训练流程相同，只需更换数据 YAML、类别和关键点定义。若使用 5 个关键点，则设置 kpt_shape: [5, 3]，同时写出与任务结构对应的 flip_idx。

7. 验证与可视化

yolo pose val \
  model=runs/armor_pose/baseline/weights/best.pt \
  data=armor-pose.yaml \
  imgsz=640 \
  device=0

yolo pose predict \
  model=runs/armor_pose/baseline/weights/best.pt \
  source=/path/to/test-video.mp4 \
  conf=0.25 \
  save=True

不要只看训练日志中的总指标。至少观看几段连续测试视频，重点观察：

远距离目标是否连续出现；
目标快速运动时角点是否落在真实灯条端点；
复杂背景是否产生持续误检；
相邻帧的类别和关键点是否抖动；
部分遮挡时，模型是在合理退化还是输出错误结构。

8. 导出 ONNX

yolo export \
  model=runs/armor_pose/baseline/weights/best.pt \
  format=onnx \
  imgsz=640 \
  opset=17 \
  simplify=True \
  dynamic=False

导出后应使用同一批图片比较 PyTorch 与 ONNX 输出。检查输入名称、形状、颜色顺序、归一化方式、输出维度和关键点坐标恢复是否一致，再进入 TensorRT 或 OpenVINO 转换。

当前 3SE 装甲板生产模型使用的是深圳大学 0526.onnx 的特定 22 维输出协议，并不是 Ultralytics Pose 的默认输出协议。自己训练的新模型不能只替换权重文件，还需要让 C++ 后处理与新模型的输出张量保持一致。

9. 常见训练问题

现象	常见原因	检查方向
训练很快收敛，测试视频很差	相邻视频帧同时进入训练和验证集	按录制场次重新划分数据
框正确但角点交叉	标注顺序或 `flip_idx` 错误	绘制增强后的标签检查语义顺序
远距离目标大量漏检	小目标样本少、输入分辨率低	增加远距离数据，比较更大 `imgsz`
场地灯误检	负样本不足	加入包含相似灯光但无目标的图像
夜间或逆光失效	训练集光照覆盖不足	补充真实数据，谨慎增加亮度和噪声增强
loss 出现 NaN	标签越界、空标签格式错误、学习率异常	扫描标签并检查首批可视化
CUDA 显存不足	batch、分辨率或模型过大	减小 batch/imgsz，先使用 n 模型
ONNX 与训练结果不同	预处理、opset、动态形状或后处理不一致	固定输入逐层比较输出
mAP 很高但 PnP 抖动	框指标掩盖了角点误差	单独统计角点误差和 PnP 位姿误差