RGB-D 分割与 3D 定位系统

基于 RGB-D 相机的实时目标分割、3D 定位和空间风险感知系统。支持 YOLO + SAM2 实例分割、点云处理、目标跟踪和风险估计。

✨ 特性

• 🎯 实时实例分割：基于 YOLOv11-seg 和 SAM2 的高精度目标分割
• 📍 3D 定位：从 RGB-D 数据反投影到 3D 点云，实现目标精确定位
• ⚠️ 风险感知：计算目标与环境点云的最近距离，评估碰撞风险
• 🎨 可视化界面：实时 GUI 预览、目标锁定和状态监控
• 📊 离线处理：批量处理已录制序列，生成可视化报告
• 🔧 可配置流水线：模块化设计，支持灵活配置和扩展

📋 目录结构

RGB-D_Segmentation_and_3D_Localization_for_Real-time_Target_Locking_and_Spatial_Risk_Awareness/
├── configs/              # 配置文件
│   ├── default.yaml      # 默认配置
│   └── postprocess/      # 后处理流水线配置
├── data/                 # 数据目录（示例数据）
│   └── sample_01/        # 示例序列
├── models/               # 模型权重目录
│   ├── yolo-seg/         # YOLO 分割模型
│   └── sam/              # SAM2 模型
├── outputs/              # 输出目录
│   └── sample_01/        # 处理结果
├── scripts/              # 脚本文件
│   ├── record_orbbec_tk.py    # GUI 采集与锁定
│   ├── test_output.py          # 离线处理与可视化
│   ├── run_realtime.py         # 实时处理
│   └── calibrate_*.py          # 标定脚本
├── src/                  # 源代码
│   ├── perception/      # 感知模块
│   ├── risk/            # 风险估计模块
│   ├── utils/           # 工具模块
│   ├── control/         # 控制模块（机器人）
│   └── calibration/     # 标定模块
├── tests/               # 测试文件
├── requirements.txt     # Python 依赖
├── environment.yml      # Conda 环境配置
└── README.md           # 本文档

🚀 快速开始

环境要求

• Python 3.10
• CUDA 12.8（可选，用于 GPU 加速）
• Orbbec RGB-D 相机（或使用离线数据）

安装步骤

1. 克隆仓库

git clone https://github.com/checheng117/CIE6004_GroupProject_Team79
cd RGB-D_Segmentation_and_3D_Localization_for_Real-time_Target_Locking_and_Spatial_Risk_Awareness

2. 创建 Conda 环境

conda env create -f environment.yml
conda activate rag

3. 安装 PyTorch（CUDA 12.8）

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

4. 安装其他依赖

pip install -r requirements.txt

5. 下载模型权重

将以下模型文件放置在对应目录：

• models/yolo-seg/yolo11s-seg.pt - YOLO11 分割模型（小）
• models/sam/sam2.1_b.pt - SAM2 基础模型

模型下载链接：

• YOLO11
• SAM2

📖 使用指南

GUI 实时采集与目标锁定

启动 GUI 界面进行实时 RGB-D 采集和目标锁定：

python scripts/record_orbbec_tk.py \
  --model models/yolo-seg/yolo11s-seg.pt \
  --sam2-model models/sam/sam2.1_b.pt \
  --seg-mode yolo+sam2 \
  --pipeline-cfg configs/postprocess/default.yaml \
  --color 1280x720@30 \
  --depth 640x576@30 \
  --out data/sample_01

操作说明：

• 点击图像中的目标对象进行锁定
• 按 R 键开始/停止录制序列
• 按 Q 或 Esc 退出

输出文件：

• data/sample_01/rgb/*.png - RGB 图像序列
• data/sample_01/depth/*.png - 深度图像序列
• data/sample_01/mask/*.png - Mask 序列
• data/sample_01/intrinsics.json - 相机内参
• data/sample_01/target_state.jsonl - 每帧目标状态

离线处理与可视化

批量处理已录制的序列，生成可视化报告：

python test_output.py \
  --input data/sample_01 \
  --outdir outputs/ppt_export \
  --frames 0 50 \
  --seg-mode yolo+sam2 \
  --model models/yolo-seg/yolo11s-seg.pt \
  --sam2-model models/sam/sam2.1_b.pt \
  --pipeline-cfg configs/postprocess/default.yaml \
  --mask-align-mode auto

输出内容：

• 每帧可视化图像（RGB、分割、3D 定位等）
• 汇总统计图表（距离、风险值趋势）
• JSON 元数据文件

自动选择最佳 Mask 对齐模式

系统支持多种 mask 对齐模式，可自动选择最佳模式：

python test_output.py \
  --input data/sample_01 \
  --mask-align-mode auto

对齐模式：

• resize - 直接缩放（默认，最快）
• center-crop - 中心裁剪（保持宽高比）
• letterbox - 填充黑边（保留全部信息）
• auto - 自动选择最佳模式

对比所有对齐模式

生成所有对齐模式的对比结果：

python test_output.py \
  --input data/sample_01 \
  --compare-align-modes

🔧 配置说明

后处理流水线配置

编辑 configs/postprocess/default.yaml 自定义流水线参数：

pipeline:
  stages:
    - stage: back_projector
      type: pointcloud_basic
      params:
        depth_min: 0.20      # 最小深度（米）
        depth_max: 1.50      # 最大深度（米）
        max_target_points: 5000
        max_env_points: 20000
    
    - stage: risk_estimator
      type: risk_estimator
      params:
        tau: 0.05            # 风险阈值（米）
        sigma: 0.05          # 风险衰减系数
        ema_alpha: 0.35      # 时序平滑系数

命令行参数

主要参数说明：

GUI 采集 (record_orbbec_tk.py):

• --model - YOLO 模型路径
• --sam2-model - SAM2 模型路径
• --seg-mode - 分割模式：yolo / sam2 / yolo+sam2
• --pipeline-cfg - 流水线配置文件
• --color / --depth - RGB/深度分辨率
• --enable-robot-send - 启用 Socket 通信

离线处理 (test_output.py):

• --input - 输入序列目录
• --outdir - 输出目录
• --frames - 处理的帧范围
• --mask-align-mode - Mask 对齐模式
• --depth-scale / --depth-unit - 深度单位设置

🏗️ 系统架构

数据流

RGB-D 输入
  ↓
[1] 图像采集与预处理
  ├─ OrbbecCamera：RGB-D 同步采集
  ├─ 深度归一化：自动推断单位 → 统一转换为米
  └─ 内参加载：从 intrinsics.json 读取
  ↓
[2] 实例分割（YOLO / SAM2）
  ├─ YOLO 检测：生成初始 mask 和 bbox
  └─ SAM2 精细化：基于 bbox 精细化 mask
  ↓
[3] Mask 对齐
  └─ 将 RGB mask 对齐到 Depth 分辨率
  ↓
[4] 后处理流水线
  ├─ 反投影：depth + mask → 3D 点云（S, O）
  ├─ 平面分割：RANSAC 拟合主平面
  ├─ 欧式聚类：将目标点云拆分为多个簇
  ├─ 簇选择：选择目标对象
  ├─ 环境构建：构建环境点云
  └─ 风险估计：计算最近距离和风险值
  ↓
[5] 目标估计
  └─ 从 mask 内点云估计目标 3D 中心
  ↓
[6] 输出与可视化
  ├─ GUI 显示
  ├─ Socket JSON 发送
  └─ 日志文件

核心模块

• 感知模块 (src/perception/): 分割、点云处理、目标估计
• 风险模块 (src/risk/): 距离计算、风险映射、时序平滑
• 工具模块 (src/utils/): 相机接口、流水线框架、可视化
• 控制模块 (src/control/): 机器人控制接口
• 标定模块 (src/calibration/): 手眼标定

📊 输出格式

目标状态 JSON

{
  "status": "locked",
  "target_id": 0,
  "conf": 0.95,
  "frame": "cam",
  "p_target_cam": [0.5, 0.2, 0.8],
  "p_target": [0.5, 0.2, 0.8],
  "d_min": 0.12,
  "risk": 0.09,
  "health": [],
  "timestamp": 1234567890.123
}

Socket 通信

系统支持通过 TCP Socket 发送目标状态（默认 15Hz）：

python scripts/record_orbbec_tk.py \
  --enable-robot-send \
  --robot-host 127.0.0.1 \
  --robot-port 8888 \
  --send-rate-hz 15.0

🧪 测试

运行测试套件：

pytest tests/

主要测试模块：

• test_back_projector.py - 反投影测试
• test_risk_estimator.py - 风险估计测试
• test_selector.py - 簇选择测试
• test_pipeline_modules.py - 流水线模块测试

🐛 故障排除

常见问题

1. MASK_EMPTY - 未检测到目标

   • 降低置信度阈值：--conf 0.15
   • 检查目标是否在视野内
   • 尝试不同的分割模式

2. 深度单位错误 - 目标位置异常

   • 使用 --depth-unit mm 或 --depth-scale 0.001
   • 检查 intrinsics.json 中的内参

3. SAM2 初始化失败

   • 检查权重文件是否存在
   • 使用 CPU 模式：--yolo-device cpu
   • 使用更小的模型：sam2.1_s.pt

4. 处理速度慢

   • 使用 GPU：--yolo-device cuda
   • 使用小模型：yolo11s-seg.pt
   • 启用 FP16：--half

更多故障排除信息请参考原 README 文档的"排障指南"部分。

📝 许可证

请查看 LICENSE 文件。

🙏 致谢

• Ultralytics YOLO - YOLO11 分割模型
• Meta SAM2 - SAM2 分割模型
• Orbbec SDK - RGB-D 相机支持

📧 联系方式

如有问题或建议，请提交 Issue 或 Pull Request。

---