# Armor YOLO Detect

基于 **YOLOv5** 和 **TensorRT** 的装甲板检测节点，用于 RoboMaster 机甲大师竞赛的自瞄系统。

本节点作为原始二值化+灯条匹配方案（`armor_detector`）的替代方案，使用深度学习模型直接检测装甲板，提供更好的泛化能力和检测精度。

---

## 目录

- [简介](#简介)
- [目录结构](#目录结构)
- [模型说明](#模型说明)
- [快速开始](#快速开始)
  - [环境依赖](#环境依赖)
  - [模型转换](#模型转换)
  - [编译项目](#编译项目)
  - [运行节点](#运行节点)
- [详细使用](#详细使用)
- [参数说明](#参数说明)
- [与原始方案对比](#与原始方案对比)
- [调试方法](#调试方法)
- [常见问题](#常见问题)

---

## 简介

本项目使用 RobotDetectionModel 中的 0526 版本 ONNX 模型，该模型：
- **网络结构**: 魔改 YOLOv5，backbone 采用 MobileNetV3
- **训练数据**: 约 15K 张高质量数据集
- **输入**: 640x640 RGB 彩色图像
- **输出**: 25200x85 矩阵，包含四个角点、置信度、颜色分类和数字分类
- **推理设备**: Jetson NX (TensorRT 加速)

---

## 目录结构

```
armor_yolo_detect/
├── include/armor_yolo_detect/
│   ├── types.hpp              # 数据结构：Armor, YoloObject, NumberClass 等
│   ├── yolo_tensorrt.hpp/cpp  # TensorRT 推理引擎封装
│   ├── trt_logger.hpp/cpp     # TensorRT 日志工具
│   ├── armor_yolo_detector.hpp/cpp  # 检测器类封装
│   └── armor_yolo_detector_node.hpp/cpp  # ROS2 节点实现
├── src/                       # 源文件
├── config/
│   └── params.yaml           # 默认参数配置
├── model/
│   └── label.yaml            # 类别标签定义
└── launch/
    └── armor_yolo_detector.launch.py  # 启动文件
```

---

## 模型说明

### 输入格式

| 属性 | 值 |
|------|-----|
| 尺寸 | 640 x 640 |
| 通道 | 3 (RGB 彩色图) |
| 数据类型 | float32 |
| 归一化 | [0, 1] |
| 格式 | NCHW (TensorRT) |

### 输出格式 (25200 x 85)

每行代表一个检测候选框，85 个通道含义如下：

| 索引 | 含义 | 说明 |
|------|------|------|
| 0-7 | 四个角点坐标 | (x1,y1,x2,y2,x3,y3,x4,y4)，左上角开始逆时针 |
| 8 | 置信度 | sigmoid 激活 |
| 9-12 | 颜色分类 | 红0, 蓝1, 灰2, 紫3 |
| 13-21 | 数字分类 | 0:G, 1:1, 2:2, 3:3, 4:4, 5:5, 6:O, 7:Bs, 8:Bb |

### 类别定义

**数字类别 (13-21):**
- `0`: G (哨兵 Guard)
- `1`: 1 (一号步兵)
- `2`: 2 (二号步兵)
- `3`: 3 (三号步兵)
- `4`: 4 (四号步兵)
- `5`: 5 (五号步兵)
- `6`: O (前哨站 Outpost)
- `7`: Bs (基地小装甲)
- `8`: Bb (基地大装甲)

**颜色类别 (9-12):**
- `0`: RED (红色敌方)
- `1`: BLUE (蓝色敌方)
- `2`: GRAY (灰色-无威胁)
- `3`: PURPLE (紫色-无威胁)

---

## 快速开始

### 环境依赖

#### Jetson NX 平台

Jetson NX 已预装 CUDA，需要安装 TensorRT：

```bash
# 1. 安装 Python TensorRT 和相关包
sudo apt update
sudo apt install python3-libnvinfer-dev tensorrt

# 2. 验证安装
python3 -c "import tensorrt; print(tensorrt.__version__)"

# 3. 确保 CUDA 环境变量
echo 'export LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
```

#### 其他依赖

```bash
# OpenCV (通常已预装)
sudo apt install libopencv-dev python3-opencv

# ROS2 Humble (参考 rm_auto_aim 原有配置)
# 确保已正确安装并 source
source /opt/ros/humble/setup.bash
```

### 模型转换

**注意**: 模型转换需要在有 NVIDIA GPU 的机器上进行（可以是 Jetson NX 本身）

```bash
# 进入工作空间
cd ~/colcon_ws

# 一键转换（推荐）
# 语法: ./scripts/build_engine.sh --onnx <onnx路径> --engine <输出engine路径> [选项]
./tools/scripts/build_engine.sh \
    --onnx ../RobotDetectionModel/Model/0526.onnx \
    --engine src/rm_auto_aim/armor_yolo_detect/model/yolov5_0526.engine \
    --fp16

# 完整参数说明:
# --onnx      输入 ONNX 模型路径
# --engine    输出 TensorRT engine 路径
# --batch-size 批处理大小，默认 1
# --workspace  工作空间大小(GB)，默认 4
# --fp16      启用 FP16 加速（Jetson NX 推荐）
# --int8      启用 INT8 加速（需要校准）
# --verbose   显示详细日志
```

**转换成功标志**:
```
Building TensorRT engine from ONNX...
[TensorRT] INFO: [MemUsageChange] Init CUDA engine for native input sizes...
Engine built successfully!
Engine saved to: /path/to/engine
Conversion completed!
```

### 编译项目

```bash
cd ~/colcon_ws

# 编译 armor_yolo_detect 包
colcon build --packages-select armor_yolo_detect

# 如果需要编译依赖
colcon build --packages-select armor_yolo_detect rm_interfaces rm_utils

# source 环境
source install/setup.bash
```

### 运行节点

```bash
# 启动 YOLO 检测节点（红色敌方）
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py detect_color:=0

# 蓝色敌方
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py detect_color:=1

# 启用调试模式
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py debug:=true
```

---

## 详细使用

### 完整启动示例

```bash
# 1. Source 环境
source /opt/ros/humble/setup.bash
source ~/colcon_ws/install/setup.bash

# 2. 启动相机节点（如有需要）
# ros2 launch hik_camera hik_camera.launch.py

# 3. 启动 YOLO 检测节点
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py \
    detect_color:=0 \
    conf_threshold:=0.7 \
    nms_threshold:=0.5 \
    debug:=true
```

### 参数配置

可以在 `config/params.yaml` 中修改默认参数：

```yaml
/**:
  ros__parameters:
    # 检测参数
    conf_threshold: 0.65      # 置信度阈值 (0.0-1.0)
    nms_threshold: 0.45       # NMS 阈值 (0.0-1.0)
    detect_color: 0          # 0=红色, 1=蓝色

    # 处理参数
    max_queue_size: 3        # 图像队列大小
    process_every_n_frames: 1  # 每N帧处理一次

    # PnP 和 BA
    use_ba: true             # 是否使用 BA 优化

    # 坐标系
    target_frame: "odom"     # TF 目标坐标系

    # 调试
    debug: false
    debug:
      enable_terminal_log: true
      enable_markers: true
      enable_result_img: true
    debug_log_interval_frames: 30
```

### 动态调参

运行时可以通过命令行调整参数：

```bash
# 降低置信度阈值（检测更多目标）
ros2 param set /armor_yolo_detector conf_threshold 0.5

# 提高 NMS 阈值（合并更近的目标）
ros2 param set /armor_yolo_detector nms_threshold 0.6

# 切换敌方颜色
ros2 service call /armor_detector/set_mode rm_interfaces/srv/SetMode "{mode: 0}"  # 红
ros2 service call /armor_detector/set_mode rm_interfaces/srv/SetMode "{mode: 1}"  # 蓝
```

---

## 参数说明

| 参数 | 类型 | 默认值 | 说明 |
|------|------|--------|------|
| `conf_threshold` | float | 0.65 | 置信度阈值，低于此值的检测框将被丢弃 |
| `nms_threshold` | float | 0.45 | NMS（非极大值抑制）阈值，用于去除重叠框 |
| `detect_color` | int | 0 | 0=红色敌方, 1=蓝色敌方 |
| `max_queue_size` | int | 3 | 图像队列最大长度 |
| `process_every_n_frames` | int | 1 | 每 N 帧处理一次，用于降低帧率 |
| `use_ba` | bool | true | 启用 BA（束调整）优化位姿 |
| `target_frame` | string | "odom" | TF 变换的目标坐标系 |
| `debug` | bool | false | 是否启用调试模式 |

### 建议参数调整

- **低光照环境**: 降低 `conf_threshold` 到 0.5-0.6
- **误检测多**: 提高 `conf_threshold` 到 0.7-0.8
- **目标频繁切换**: 提高 `nms_threshold` 到 0.5-0.6

---

## 与原始方案对比

| 特性 | armor_detector (原始) | armor_yolo_detect (本项目) |
|------|----------------------|---------------------------|
| 检测方式 | 二值化 + 灯条匹配 | YOLOv5 深度学习 |
| 输入 | 灰度图 | RGB 彩色图 |
| 泛化能力 | 依赖阈值参数 | 模型自适应 |
| 计算量 | 较低 | 较高 |
| 依赖 | OpenCV | TensorRT + CUDA |
| 输出话题 | `/armor_detector/armors` | `/armor_detector/armors` |
| 消息类型 | `rm_interfaces::msg::Armors` | `rm_interfaces::msg::Armors` |

**切换方案**:
```bash
# 原始二值化方案
ros2 launch rm_bringup bringup.launch.py detector:=binary

# YOLO 方案
ros2 launch rm_bringup bringup.launch.py detector:=yolo
```

---

## 调试方法

### 启用调试输出

```bash
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py debug:=true
```

### 查看话题

```bash
# 查看检测结果图像
ros2 topic echo /armor_yolo_detector/result_img

# 查看原始图像（需要 remap）
ros2 run image_view image_view --image /armor_yolo_detector/result_img

# 查看检测到的装甲板
ros2 topic echo /armor_detector/armors

# 查看可视化标记
ros2 topic echo /armor_detector/marker
```

### Rviz2 可视化

```bash
# 启动 rviz2
rviz2

# 添加 Display:
# - Image: /armor_yolo_detector/result_img
# - Marker: /armor_detector/marker
```

### 发表话题重映射

如果相机话题名称不同，需要重映射：

```bash
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py \
    remappings:="[('/image_raw', '/your_camera/image_raw'), ('/camera_info', '/your_camera/camera_info')]"
```

---

## 常见问题

### Q1: 编译报错 "TensorRT not found"

```bash
# 确保安装了 TensorRT 开发包
sudo apt install tensorrt libnvinfer-dev python3-libnvinfer-dev

# 验证
dpkg -l | grep TensorRT
```

### Q2: 运行时报错 "Engine file not found"

```bash
# 检查 engine 文件是否存在
ls -la ~/colcon_ws/src/rm_auto_aim/armor_yolo_detect/model/

# 如果没有，执行模型转换
./tools/scripts/build_engine.sh --onnx ../RobotDetectionModel/Model/0526.onnx --engine src/rm_auto_aim/armor_yolo_detect/model/yolov5_0526.engine --fp16
```

### Q3: 检测框错位或翻转

检查图像输入格式，确保是 RGB 输入：
```bash
# 查看图像编码
ros2 topic echo /hik_camera/image_raw --field encoding
```
如果是 `mono8`（灰度），需要修改代码中的预处理逻辑。

### Q4: Jetson NX 上推理速度慢

```bash
# 1. 确保使用 FP16
./tools/scripts/build_engine.sh --onnx model.onnx --engine model.engine --fp16

# 2. 检查电源模式（JetPack 4.6+）
sudo nvpmodel -m 0  # 最大性能模式

# 3. 查看推理时间
ros2 run rqt_top rqt_top
```

### Q5: 如何同时运行多个检测器进行对比?

```bash
# 终端 1: 运行原始方案
ros2 launch armor_detector armor_detector.launch.py

# 终端 2: 运行 YOLO 方案（修改节点名避免冲突）
ros2 launch armor_yolo_detect armor_yolo_detector.launch.py \
    --remap /armor_detector/armors:=/armor_detector/armors_yolo
```

---

## 维护者

- **开发者**: chenyy
- **邮箱**: 3289288508@qq.com
- **开源**: https://github.com/chenjunnn/rm_auto_aim

## 协议

本项目基于 Apache 2.0 协议开源。
