errm / docs /CODE_GUIDE.md

yuffish

Add files using upload-large-folder tool

a741a7c verified 8 days ago

preview code

raw

history blame contribute delete

12.5 kB

代码指南 - Code Guide

详细说明项目中每个脚本的功能、输入输出和使用方法

📖 目录

数据提取
视频标注
数据转换
数据统计
视频描述生成
工具脚本

1. 数据提取

video_process.py

功能：从 DROID TFDS 数据集中提取失败的 episode 并保存为视频和元数据。

输入：

DROID TFDS 数据集（/playpen-ssd/dataset/）

输出：

视频文件：./droid_processed/episode_XXXXXX.mp4
元数据文件：./droid_processed/episode_XXXXXX.json

主要功能：

判断 episode 是否失败（is_last=True 且 is_terminal=False）
拼接 wrist_image_left 和 exterior_image_1_left 两个视角
保存视频（15fps，h264 编码）
保存元数据（包含语言指令、动作序列等）

使用方法：

cd /home/jqliu/projects/RewardModel
python video_process.py

配置：

FAILED_DIR：输出目录
数据集路径：/playpen-ssd/dataset/

video_process_v2.py

功能：改进版数据提取脚本（功能类似 video_process.py）。

特点：

优化了失败判断逻辑
改进了元数据存储格式

extract_frames_to_images.py

功能：从标注结果（JSONL）中提取对应的视频帧，保存为图片。

输入：

标注文件：./output/labels_batch.jsonl
原始视频文件（从标注中读取路径）

输出：

图片文件：./data/frames/<date>/<session>/<video_id>_frame_<idx>.jpg
索引文件：./data/frames/frame_index_map.csv

保存模式：

mirror（镜像模式）：按原始目录结构保存

./data/frames/2023-07-07/session_xxx/22008760_frame_0.jpg

flat_hash（扁平哈希模式）：所有图片保存在同一目录，文件名包含路径哈希
```
./data/frames/22008760__2023-07-07__session_xxx__abc12345_f0.jpg
```

使用方法：

python extract_frames_to_images.py

配置：

jsonl_path = "./output/labels_batch.jsonl"
output_image_root = Path("./data/frames")
SAVE_MODE = "mirror"  # 或 "flat_hash"
write_index = True    # 是否生成索引文件

2. 视频标注

api_batch.py

功能：基础版 GPT 视频标注脚本。

特点：

一次性发送所有帧给 GPT
基础奖励指标

不推荐使用，建议使用改进版。

api_batch_improved.py ⭐ 推荐

功能：改进版 GPT 视频标注脚本，支持滑动窗口和动态任务描述。

输入：

视频文件（自动搜索指定目录）
Metadata JSON 文件（自动查找）

输出：

标注文件：./output/labels_batch_improved.jsonl

主要特性：

动态任务描述
- 自动从 metadata JSON 提取 current_task 字段
- 动态插入到 GPT prompt 中
滑动窗口分析（推荐开启）
- 连续帧分析，捕捉时序动态
- 上下文传递（前一窗口结果作为下一窗口的上下文）
- 参数：WINDOW_SIZE=5, WINDOW_STRIDE=3
奖励指标对齐
- 与仿真数据对齐：reachout, grasp, collision, fall, smooth
- 6 阶段定义：reach, grasp, lift, move, place, retract

使用方法：

cd /home/jqliu/projects/RewardModel

# 设置 API Key
export OPENAI_API_KEY="your-key-here"

# 运行标注
python api_batch_improved.py

配置：

ROOT = Path("/playpen-ssd/dataset/droid_raw/1.0.1/AUTOLab/failure")
TARGET_NAME = "22008760.mp4"   # 视频文件名模式
OUTPUT_FILE = "./output/labels_batch_improved.jsonl"
FPS_SAMPLE = 2                 # 帧采样率
MODEL_NAME = "gpt-4o"          # GPT 模型
MAX_VIDEOS = 10                # 最多处理多少个视频
START_INDEX = 0                # 从第几个开始

# 滑动窗口配置
USE_SLIDING_WINDOW = True      # 是否使用滑动窗口
WINDOW_SIZE = 5                # 窗口大小
WINDOW_STRIDE = 3              # 滑动步长

输出格式：

{
  "video_path": "/path/to/video.mp4",
  "video_id": "22008760",
  "task": "Move object into container",
  "t": 0,
  "frame_index": 0,
  "total_frames": 229,
  "native_fps": 60.0,
  "window_idx": 0,
  "stage": 0,
  "stage_name": "reach",
  "reachout": 0.3,
  "grasp": 0.0,
  "collision": 0.0,
  "fall": 0.0,
  "smooth": 0.8,
  "reward": 0.55,
  "delta": 0.0,
  "success_prob": 0.5,
  "failure": 0,
  "explanation": "Robot arm is reaching toward the container..."
}

详细说明：参见 README_improved_api.md

api_batch_improved_v2.py

功能：使用 Claude 模型的视频标注脚本（而非 GPT）。

特点：

使用本地部署的 Claude API（http://localhost:9950/v1）
功能与 api_batch_improved.py 类似
适合有本地 Claude 服务的场景

配置：

client = OpenAI(
    base_url="http://localhost:9950/v1",
    api_key="your-api-key"
)
MODEL_NAME = "anthropic.claude-3-7-sonnet-20250219-v1:0"

api_batch_hf_dataset.py

功能：处理 HuggingFace 数据集格式的视频标注。

输入：

HuggingFace 数据集目录：/playpen-ssd/dataset/droid_raw/hg_data
数据集信息文件：dataset_info.json

输出：

标注文件：./output/labels_hf_dataset.jsonl

主要特性：

读取 HuggingFace 数据集元数据
支持批量处理
与改进版 API 相同的标注逻辑

使用方法：

python api_batch_hf_dataset.py

配置：

DATASET_ROOT = Path("/playpen-ssd/dataset/droid_raw/hg_data")
OUTPUT_FILE = "./output/labels_hf_dataset.jsonl"
START_INDEX = 0
MAX_VIDEOS = 10
USE_SLIDING_WINDOW = True

3. 数据转换

convert_to_sft.py

功能：将标注结果转换为 SFT（Supervised Fine-Tuning）训练格式。

输入：

标注文件：./output/labels_batch.jsonl
帧图片：./data/frames/

输出：

中间格式：./data/sft_frames_dataset.json

{
  "conversation": [
    {
      "role": "user",
      "content": [
        {"type": "image", "image": "/path/to/frame.jpg"},
        {"type": "text", "text": "Please output the execution stage and reward..."}
      ]
    },
    {
      "role": "assistant",
      "content": "{\"stage\": 0, \"reward\": 0.5, ...}"
    }
  ]
}

最终格式：./data/sft_frames_dataset_sft.json

{
  "conversation": [
    {"from": "human", "value": "<Image: /path/to/frame.jpg>\nPlease output..."},
    {"from": "assistant", "value": "{\"stage\": 0, \"reward\": 0.5, ...}"}
  ]
}

使用方法：

python convert_to_sft.py

注意事项：

确保 extract_frames_to_images.py 已运行，图片已提取
确保帧图片路径与标注中的路径匹配

4. 数据统计

位置：./data_sta/

count_failure_trajectories.py

功能：统计 DROID 数据集中的失败轨迹。

输入：

DROID 原始数据：/playpen-ssd/dataset/droid_raw/1.0.1/

输出：

JSON 格式：./data_sta/failure_statistics.json
文本格式：./data_sta/failure_statistics.txt

统计内容：

总轨迹数
失败轨迹数
按数据源分布
按任务类型分布
Top 任务类型排名

使用方法：

cd /home/jqliu/projects/RewardModel/data_sta
python count_failure_trajectories.py

check_video_framerate.py

功能：检查 DROID 数据集中视频的帧率统计。

输出：

./data_sta/video_framerate_stats.json
./data_sta/video_framerate_stats.txt

统计内容：

视频帧率分布
视频时长统计
帧数统计

prepare_hf_dataset.py

功能：从 DROID 数据集采样视频，准备上传到 HuggingFace。

主要功能：

统计可用视频和元数据
按策略采样视频
复制到输出目录
生成 dataset_info.json 元数据

采样策略：

balanced：每个任务类别采样相同数量
random：完全随机采样
proportional：按原始任务分布比例采样

使用方法：

cd /home/jqliu/projects/RewardModel/data_sta
python prepare_hf_dataset.py

配置：

TOTAL_SAMPLES = 2500           # 总采样数
PROCESS_TYPE = "failure"       # failure/success/both
SAMPLING_STRATEGY = "balanced" # 采样策略
VIDEO_PATTERN = "2*.mp4"       # 视频文件名模式
EXCLUDE_STEREO = True          # 排除 stereo 视频

详细说明：参见 data_sta/README_HF_DATASET.md

upload_to_huggingface.py

功能：上传准备好的数据集到 HuggingFace Hub。

输入：

准备好的数据集目录：/playpen-ssd/dataset/droid_raw/hg_data

使用方法：

# 设置 HuggingFace Token
export HF_TOKEN="your_token_here"

# 上传
python upload_to_huggingface.py

配置：

HF_USERNAME = "your-username"
DATASET_NAME = "droid-failure-sampled"
PRIVATE = False  # True=私有, False=公开

详细说明：参见 data_sta/UPLOAD_GUIDE.md

5. 视频描述生成

位置：./caption/

infer_caption.py

功能：使用 Qwen-VL 模型生成单个视频的描述。

输入：

视频文件

输出：

视频描述文本（JSON 格式）

infer_caption_batch.py

功能：批量生成视频描述。

输入：

多个视频文件

输出：

./caption/output/caption_results_batch.json

主要特性：

批量处理多个视频
使用 Qwen-VL 模型
JSON 自动解析

使用方法：

cd /home/jqliu/projects/RewardModel/caption
python infer_caption_batch.py

详细说明：参见 caption/README_batch.md

6. 工具脚本

api.py

功能：早期的单视频 API 调用脚本（已弃用）。

建议：使用 api_batch_improved.py 代替。

📚 脚本依赖关系

数据提取流程：
1. video_process.py → droid_processed/*.mp4, *.json

标注和训练数据准备流程：
2. api_batch_improved.py → output/labels_batch_improved.jsonl
3. extract_frames_to_images.py → data/frames/**/*.jpg
4. convert_to_sft.py → data/sft_frames_dataset_sft.json

HuggingFace 数据集流程：
5. count_failure_trajectories.py → failure_statistics.json
6. prepare_hf_dataset.py → hg_data/
7. upload_to_huggingface.py → HuggingFace Hub

标注数据处理流程（HF）：
8. api_batch_hf_dataset.py → output/labels_hf_dataset.jsonl

视频描述流程：
9. infer_caption_batch.py → caption/output/caption_results_batch.json

🔍 常见使用场景

场景 1：处理新的视频数据

# 1. 使用改进版 API 标注视频
python api_batch_improved.py

# 2. 提取帧图片
python extract_frames_to_images.py

# 3. 转换为训练格式
python convert_to_sft.py

场景 2：准备和上传 HuggingFace 数据集

cd data_sta

# 1. 统计数据
python count_failure_trajectories.py

# 2. 采样和准备数据集
python prepare_hf_dataset.py

# 3. 上传到 HuggingFace
export HF_TOKEN="your-token"
python upload_to_huggingface.py

场景 3：处理 HuggingFace 数据集

# 1. 标注 HuggingFace 数据集
python api_batch_hf_dataset.py

# 2. 提取帧和转换格式（与场景1相同）
python extract_frames_to_images.py
python convert_to_sft.py

⚠️ 注意事项

API Key 配置
- GPT 相关脚本需要设置 OPENAI_API_KEY
- HuggingFace 上传需要设置 HF_TOKEN
路径配置
- 确保数据集路径正确
- 相对路径基于项目根目录

依赖安装

pip install openai decord pillow python-dotenv tqdm tensorflow_datasets

数据量控制
- 首次运行建议设置 MAX_VIDEOS 为小值（如 1-10）测试
- 成功后再扩大规模
错误处理
- 所有脚本都有 try-except 错误处理
- 失败的视频会跳过并记录日志
- 可以重新运行继续处理

yuffish
/

errm

代码指南 - Code Guide

📖 目录

1. 数据提取

video_process.py

video_process_v2.py

extract_frames_to_images.py

2. 视频标注

api_batch.py

api_batch_improved.py ⭐ 推荐

api_batch_improved_v2.py

api_batch_hf_dataset.py

3. 数据转换

convert_to_sft.py

4. 数据统计

count_failure_trajectories.py

check_video_framerate.py

prepare_hf_dataset.py

upload_to_huggingface.py

5. 视频描述生成

infer_caption.py

infer_caption_batch.py

6. 工具脚本

api.py

📚 脚本依赖关系

🔍 常见使用场景

场景 1：处理新的视频数据

场景 2：准备和上传 HuggingFace 数据集

场景 3：处理 HuggingFace 数据集

⚠️ 注意事项

📖 更多文档