Dataset Tools

ROS 2 数据集采集与转换工具，用于 LeRobot v3 数据集格式。

概述

本包提供以下功能：

• Episode 录制: 通过 Action Server 控制的分段录制
• Bag 转 LeRobot: 将 ROS 2 bag 转换为 LeRobot v3 数据集格式
• 相机标定辅助: camera_alignment 支持真机 OpenCV 相机与仿真 ROS 2 image topic；仿真模式通过 rclpy 订阅图像，并可与 sim_models 中的相机调节/ArUco 标定板工具联动

仿真相关能力是可选运行路径：常规录制与 bag 转换仍只依赖 robot_config 契约和 ROS 2 bag；只有使用 camera_alignment --use_sim 处理仿真相机 topic 时，才会进入 dataset_tools → sim_models → robot_config 的跨包交互。

架构设计

仿真相机标定边界

camera_alignment --use_sim 是仿真相机标定辅助路径，不参与 episode 录制或 bag 转换主流程。该路径的职责边界如下：

1. dataset_tools.camera_alignment 负责交互式对齐 UI、参考图/参考 marker 数据保存，以及订阅仿真相机 ROS 2 image topic；
2. sim_models 负责仿真侧相机调节辅助进程、Gazebo ArUco 标定板 spawn/despawn，以及将调节结果保存为 ~/.ros/ibrobot/sim_camera_overrides/<camera>.yaml；
3. robot_config 在后续 robot.launch.py use_sim:=true 启动时读取 robot YAML 和可选 override，生成对应仿真相机位姿。

因此，sim_models 只属于 camera_alignment --use_sim 的可选运行时依赖；真机相机对齐、录制服务和 bag 转换不依赖该包。

单一真理来源 (Single Source of Truth)

所有数据集工具使用 robot_config 包下的配置文件作为唯一配置来源，例如：

src/robot_config/config/robots/so101_single_arm.yaml
├── contract.observations    ← 观测定义（相机、状态等）
├── contract.actions         ← 动作定义（arm、gripper）
├── contract.rate_hz         ← 采样率
└── control_modes            ← 运行时控制模式配置

这确保了：

• 训练数据导出与在线推理配置一致
• 无需维护重复的 contract 文件
• 配置变更自动传播到所有组件

工具

1. record_cli - 交互式录制客户端

用于控制 episode 录制的命令行工具。

启动录制服务（Ubuntu 录制服务器）：

ros2 launch robot_config robot.launch.py \
    robot_config:=so101_single_arm \
    control_mode:=teleop \
    record:=true \
    record_mode:=episodic \
    use_sim:=false

如需启用 Rerun 可视化：

ros2 launch robot_config robot.launch.py \
    robot_config:=so101_single_arm \
    control_mode:=teleop \
    record:=true \
    record_mode:=episodic \
    record_visualizer:=rerun \
    use_sim:=false

启动录制客户端（同机或另一台设置了相同 ROS_DOMAIN_ID 的机器）：

ros2 run dataset_tools record_cli

episodic 录制目录现在按 dataset 组织：

<bag_base_dir>/
└── <dataset_name>/
    ├── dataset.yaml
    └── episodes/
        ├── episode_000001/
        │   ├── metadata.yaml
        │   └── *.mcap
        └── episode_000002/
            ├── metadata.yaml
            └── *.mcap

• bag_base_dir 来自 robot_config.recording.bag_base_dir
• dataset_name 默认取 recording.dataset_name，未配置时回退到机器人名
• dataset.yaml 保存 dataset 级元信息；可选通过 recording.default_task、recording.task_family 预填任务语义
• episode 级 prompt 仍写入各自 bag 的 metadata.yaml

使用方式：

========================================================
Dataset Collection CLI
Enter prompt text to start recording. (Press Enter to reuse: 'get')
Type 'q' or 'quit' to exit.
========================================================
Prompt > get        # 输入任务描述开始录制
[INFO] 🔴 RECORDING STARTED. (Press Enter to stop early)
[INFO] ✅ RECORDING SAVED: Wrote 1894 messages to /path/to/episode
Prompt > q          # 退出

record_cli 默认按 control_mode:=teleop 工作，不触发推理侧 reset。录制模型推理过程时，将客户端控制模式设为 model_inference：

ros2 run dataset_tools record_cli --ros-args -p control_mode:=model_inference

此时每个 episode 开始前会优先调用 /action_dispatcher/reset 清理动作队列，并由 action_dispatcher best-effort 触发推理侧 policy 状态重置。可通过 reset_before_episode、dispatcher_reset_service、policy_reset_service 和 reset_timeout_sec 参数覆写对应行为、服务名和等待时间。

录制完成后，推荐直接把整个 dataset 根目录转换成 LeRobot v3 数据集：

ros2 run dataset_tools bag_to_lerobot \
    --bags-dir ~/rosbag/episodes/so101_single_arm \
    --robot-config src/robot_config/config/robots/so101_single_arm.yaml \
    --out /path/to/output_dataset

2. bag_to_lerobot - Bag 转 LeRobot 数据集

将 ROS 2 episodic dataset 根目录转换为 LeRobot v3 数据集格式。

基本用法：

ros2 run dataset_tools bag_to_lerobot \
    --bags-dir ~/rosbag/episodes/so101_single_arm \
    --robot-config src/robot_config/config/robots/so101_single_arm.yaml \
    --out /path/to/output_dataset

参数说明：

参数	说明	默认值
--bags-dir	dataset 根目录或 episodes 目录，自动发现多个 episode bag	必需
--robot-config	robot_config.yaml 路径	必需
--out	输出数据集目录	必需
--repo-id	数据集 repo_id	rosbag_v30
--no-videos	存储 PNG 图像而非视频	false
--timestamp	时间戳来源 (contract/bag/header)	contract
--image-threads	图像写入线程数	4
--chunk-size	每个 chunk 的帧数	1000

输出结构：

output_dataset/
├── videos/
│   ├── observation.images.front/
│   │   └── chunk-000/file-000.mp4
│   ├── observation.images.top/
│   └── observation.images.wrist/
├── data/
│   └── chunk-000/file-000.parquet
└── meta/
    ├── info.json
    ├── tasks.parquet
    ├── stats.json
    └── episodes/

3. episode_recorder - 录制服务节点

由 launch 文件自动启动的录制服务，提供 record_episode Action Server。

通常不需要直接运行，由 robot.launch.py 根据 record_mode:=episodic 参数自动加载。 录制结果会写到 <bag_base_dir>/<dataset_name>/episodes/episode_XXXXXX/，并在 dataset 根目录生成 dataset.yaml。

4. camera_alignment - 基于 ArUco 的相机对齐工具

用于在数据采集或复现前对齐摄像头视角。--camera-source 同时支持真机视频设备和仿真 ROS 2 image topic：

输入形态	路径	行为
/dev/videoN / 整数 / 本地视频文件	真机路径（OpenCV）	保持原有工作流，支持显式请求分辨率、帧率和采集格式
/camera/<name>/image_raw	仿真路径（ROS 2 topic）	通过 rclpy 订阅图像，并接入 sim calibration 辅助能力

真机用法：

ros2 run dataset_tools camera_alignment \
    --camera-source /dev/video0 \
    --width 640 \
    --height 480 \
    --fps 60 \
    --format MJPG \
    --reference-path /tmp/camera_reference_multi.json \
    --reference-image-path /tmp/reference_img.png

仿真用法：

ros2 run dataset_tools camera_alignment \
    --camera-source /camera/top/image_raw \
    --markerless \
    --reference-image-path ref_img/reftop.png

工具支持：

• 保存当前 ArUco 角点作为参考基准
• 实时显示与参考画面的平均像素误差
• 进入“虚影对齐”界面辅助恢复视角
• 显式请求 OpenCV 采集分辨率、帧率和采集格式，并提示实际生效值

仿真模式下额外行为：

• 启动时尝试通过 sim_models.aruco_spawner 向 Gazebo 动态插入 ArUco A4 标定板，退出时自动清理。
• 当输入是 /camera/<name>/image_raw 时，会自动切到 /camera_align/<name>/image_raw 代理话题，便于与 sim_camera_adjuster 联动。
• 按 p 保存当前帧到 capture_YYYYMMDD_HHMMSS.png。
• 按 s 时，camera_alignment 会写入 ~/.ros/ibrobot/sim_camera_overrides/<camera>.yaml 的 stub 文件；真正的位姿值由 sim_models.sim_camera_adjuster 保存并在下次仿真启动时被 robot_config 读取。

兼容性边界：

• 真机路径的 OpenCVFrameSource 工作流保持兼容；ROS 2 依赖只会在仿真 topic 路径下触发。
• markerless 模式只跳过 ArUco 误差计算，不影响真机模式的原有 reference JSON / image 行为。

详细说明见：

• docs/tools/camera_alignment.md

5. camera_isp_calibrator - 基于参考图的相机色彩对齐工具

让一台 USB 摄像头（usb_cam 节点）的画面在曝光、白平衡、增益、对比度等 方面尽可能接近一张参考图片，并把结果保存为 override JSON，下次启动 robot.launch.py 时自动复用，不修改 YAML SSOT。

运行模式：

1. ROS bridge 模式：不传 --camera_index 时，沿用原有行为， 订阅 usb_cam 节点（节点名形如 /top_camera）并通过 ROS 参数服务兜底写入；
2. OpenCV 直连模式：传入 --camera_index 时，直接从本机摄像头索引 或 /dev/video* 读取画面，不需要启动 robot.launch.py 或 usb_cam 节点。

两种模式都需要准备一张参考图片或视频（视频会取首帧）。直连模式会优先通过 v4l2-ctl 写入设备 ISP 参数；建议仍传 --camera top / --camera wrist 等名称， 这样保存的 override 文件能继续被 robot.launch.py 自动复用。

基本用法：

# 方式 A：ROS bridge 模式，先启动机器人 / 摄像头
source .shrc_local
ros2 launch robot_config robot.launch.py robot_config:=so101_single_arm control_mode:=teleop

# 另一个终端运行校准工具
source .shrc_local
ros2 run dataset_tools camera_isp_calibrator \
    --camera top \
    --reference /path/to/reference.png

# 方式 B：OpenCV 直连模式，不需要启动 robot.launch.py
source .shrc_local
ros2 run dataset_tools camera_isp_calibrator \
    --camera top \
    --camera_index /dev/video0 \
    --reference /path/to/reference.png

# 也可以使用整数索引；未传 --camera 时保存名会自动派生为 video0
ros2 run dataset_tools camera_isp_calibrator \
    --camera_index 0 \
    --reference /path/to/reference.png

界面交互（单窗口 cv2 GUI，傻瓜操作）：

• a：自动模式（Lab + Planckian 投影），4 次迭代收敛后自动停下
• c：统一 K/C/Sat 搜索（实验性，详见 §5.1）。有 ROI pair 时走 m 模式 cost，没有就走 AUTO ref-cluster cost；找不到改进会自动回退到 seed。
• 拖动滑条：手动微调 exposure / wb_kelvin / gain / brightness / contrast / saturation / sharpness（松手 0.4s 后才下发，避免抖动）
• s：保存到 ~/.ros/ibrobot/camera_isp_overrides/{camera}.json
• r：恢复启动时快照的初始值
• p：保存当前 live 帧 PNG 到工作目录
• ? / h：显示 keybinding 提示
• q：退出（有未保存改动会先警告一次，再按 q 才退出）

算法：

模式	触发	说明
Auto	按 a	sRGB→Lab 计算 P50 亮度匹配曝光；CIE xy 色度通过 McCamy 公式投影到 Planckian locus，按 delta-form 调节 white_balance kelvin。每帧迭代再读、最多 4 次。
手动滑条	拖动	直接下发 V4L2 参数到 usb_cam 节点（已强制 auto_white_balance=false / autoexposure=false）。

保存生效：保存后下次 robot.launch.py 启动时，perception.py 会自动 读取 override 并覆盖 YAML 默认值；删除 JSON 即可回退。

详细算法说明：见 临时/camera_isp_plan.md §Phase 2。

5.1 统一 K/C/Sat 色彩搜索（实验性，独立模块）

模块 dataset_tools/camera_isp/color_search.py 实现了 临时/camera_isp_unified_color_search_plan.md v4 的统一搜索路径， 与既有 4 阶段流水线（曝光/增益/亮度/锐度）并行存在，不修改任何曝光相关代码。 旧 solver / hw_pipeline 全部保留作为初值估计器与失败回退。

公共接口（pure-numpy + scipy；无 ROS / cv2 依赖）：

from dataset_tools.camera_isp.color_search import (
    KCS, SettleConfig, SearchConfig, ClusterConfig,
    kmeans_signature_lab,     # Lab 单边聚类签名
    nn_match_signatures,      # 匈牙利 ΔE2000 指派
    delta_e2000,              # CIEDE2000 (vectorised)
    quantile_distance_L,      # L* 分位数 L1
    cost_24card,              # 24 色卡 cost 工厂
    cost_ref_cluster,         # AUTO ref cost 工厂
    cost_manual_roi,          # m / ROI cost 工厂（带正则）
    frame_capture,            # settle + drop + trimmed-mean
    search_KCS,               # 主搜索 driver（直接 3D 网格 + 可选精修）
    OfflineTables,            # JSON 配置加载
)

设计原则（开放给后续迭代）：

• 三模式同构：search_KCS 接收任意 cost_fn，driver 不感知模式。
• 失败安全：未找到改进时回退到 seed 并把硬件值写回 seed。
• 可注入边界：HwWriter / FrameGrabber 协议 + sleeper 钩子让单元测试无需真实相机。
• 离线表外置：camera_isp_offline_tables.json（per device 可覆盖）承载 K/C/Sat 曲线、settle、search 参数；不再硬编码。

测试：test/test_camera_isp_color_search.py（16 个用例，覆盖 ΔE2000、聚类、匈牙利、settle、driver fallback、device caps 裁剪）。

数据流

┌─────────────────────────────────────────────────────────────┐
│   src/robot_config/config/robots/so101_single_arm.yaml     │
│  ┌─────────────────────────────────────────────────────┐   │
│  │ contract (单一真理来源)                               │   │
│  │ - observations (front, top, wrist, state)           │   │
│  │ - actions (arm, gripper)                            │   │
│  │ - rate_hz: 20                                       │   │
│  └─────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────┘
                              │
          ┌───────────────────┼───────────────────┐
          ▼                   ▼                   ▼
   ┌─────────────┐     ┌─────────────┐     ┌─────────────┐
   │  录制服务    │     │  数据转换    │     │  推理服务    │
   │ episode_    │     │ bag_to_     │     │ lerobot_    │
   │ recorder    │     │ lerobot     │     │ policy_node │
   └─────────────┘     └─────────────┘     └─────────────┘
          │                   │                   │
          ▼                   ▼                   ▼
   ROS 2 Bag          LeRobot Dataset      Model Inference

配置示例

src/robot_config/config/robots/so101_single_arm.yaml 中的 contract 配置：

robot:
  name: so101_single_arm
  
  contract:
    rate_hz: 20
    max_duration_s: 90.0
    
    observations:
      - key: observation.images.front
        topic: /camera/front/image_raw
        type: sensor_msgs/msg/Image
        image:
          resize: [480, 640]
          
      - key: observation.images.top
        topic: /camera/top/image_raw
        type: sensor_msgs/msg/Image
        image:
          resize: [480, 640]
          
      - key: observation.images.wrist
        topic: /camera/wrist/image_raw
        type: sensor_msgs/msg/Image
        image:
          resize: [480, 640]
          
      - key: observation.state
        topic: /joint_states
        type: sensor_msgs/msg/JointState
        selector:
          names: [position.1, position.2, position.3, position.4, position.5, position.6]
    
    actions:
      # Arm joints (1-5)
      - key: action
        selector:
          names: [action.0, action.1, action.2, action.3, action.4]
        publish:
          topic: /arm_position_controller/commands
          type: std_msgs/msg/Float64MultiArray
          
      # Gripper joint (6) - same key for consolidation
      - key: action
        selector:
          names: [action.5]
        publish:
          topic: /gripper_position_controller/commands
          type: std_msgs/msg/Float64MultiArray

注意事项

1. Action 合并: 多个 action spec 使用相同的 key: action 会被自动合并为一个 6-DOF action
2. 观测过滤: 推理服务会根据模型的 config.json 自动过滤需要的观测
3. 录制模式:
   • record_mode:=continuous - 持续录制到一个文件
   • record_mode:=episodic - 分段录制，需要 record_cli 控制