docs/CODE_GUIDE.md

# 代码指南 - Code Guide

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

---

## 📖 目录

1. [数据提取](#1-数据提取)
2. [视频标注](#2-视频标注)
3. [数据转换](#3-数据转换)
4. [数据统计](#4-数据统计)
5. [视频描述生成](#5-视频描述生成)
6. [工具脚本](#6-工具脚本)

---

## 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 编码）
- 保存元数据（包含语言指令、动作序列等）

**使用方法**：
```bash
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`

**保存模式**：
1. **mirror**（镜像模式）：按原始目录结构保存
   ```
   ./data/frames/2023-07-07/session_xxx/22008760_frame_0.jpg
   ```

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

**使用方法**：
```bash
python extract_frames_to_images.py
```

**配置**：
```python
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`

**主要特性**：

1. **动态任务描述**
   - 自动从 metadata JSON 提取 `current_task` 字段
   - 动态插入到 GPT prompt 中

2. **滑动窗口分析**（推荐开启）
   - 连续帧分析，捕捉时序动态
   - 上下文传递（前一窗口结果作为下一窗口的上下文）
   - 参数：`WINDOW_SIZE=5`, `WINDOW_STRIDE=3`

3. **奖励指标对齐**
   - 与仿真数据对齐：reachout, grasp, collision, fall, smooth
   - 6 阶段定义：reach, grasp, lift, move, place, retract

**使用方法**：
```bash
cd /home/jqliu/projects/RewardModel

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

# 运行标注
python api_batch_improved.py
```

**配置**：
```python
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              # 滑动步长
```

**输出格式**：
```json
{
  "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](../README_improved_api.md)

---

### api_batch_improved_v2.py

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

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

**配置**：
```python
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 相同的标注逻辑

**使用方法**：
```bash
python api_batch_hf_dataset.py
```

**配置**：
```python
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/`

**输出**：
1. **中间格式**：`./data/sft_frames_dataset.json`
   ```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, ...}"
       }
     ]
   }
   ```

2. **最终格式**：`./data/sft_frames_dataset_sft.json`
   ```json
   {
     "conversation": [
       {"from": "human", "value": "<Image: /path/to/frame.jpg>\nPlease output..."},
       {"from": "assistant", "value": "{\"stage\": 0, \"reward\": 0.5, ...}"}
     ]
   }
   ```

**使用方法**：
```bash
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 任务类型排名

**使用方法**：
```bash
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。

**主要功能**：
1. 统计可用视频和元数据
2. 按策略采样视频
3. 复制到输出目录
4. 生成 `dataset_info.json` 元数据

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

**使用方法**：
```bash
cd /home/jqliu/projects/RewardModel/data_sta
python prepare_hf_dataset.py
```

**配置**：
```python
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](../data_sta/README_HF_DATASET.md)

---

### upload_to_huggingface.py

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

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

**使用方法**：
```bash
# 设置 HuggingFace Token
export HF_TOKEN="your_token_here"

# 上传
python upload_to_huggingface.py
```

**配置**：
```python
HF_USERNAME = "your-username"
DATASET_NAME = "droid-failure-sampled"
PRIVATE = False  # True=私有, False=公开
```

**详细说明**：参见 [data_sta/UPLOAD_GUIDE.md](../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 自动解析

**使用方法**：
```bash
cd /home/jqliu/projects/RewardModel/caption
python infer_caption_batch.py
```

**详细说明**：参见 [caption/README_batch.md](../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：处理新的视频数据

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

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

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

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

```bash
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 数据集

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

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

---

## ⚠️ 注意事项

1. **API Key 配置**
   - GPT 相关脚本需要设置 `OPENAI_API_KEY`
   - HuggingFace 上传需要设置 `HF_TOKEN`

2. **路径配置**
   - 确保数据集路径正确
   - 相对路径基于项目根目录

3. **依赖安装**
   ```bash
   pip install openai decord pillow python-dotenv tqdm tensorflow_datasets
   ```

4. **数据量控制**
   - 首次运行建议设置 `MAX_VIDEOS` 为小值（如 1-10）测试
   - 成功后再扩大规模

5. **错误处理**
   - 所有脚本都有 try-except 错误处理
   - 失败的视频会跳过并记录日志
   - 可以重新运行继续处理

---

## 📖 更多文档

- [项目概述](./PROJECT_OVERVIEW.md)
- [当前进展](./CURRENT_PROGRESS.md)
- [下一步计划](./NEXT_STEPS.md)
- [快速开始（AI Agent）](./QUICK_START_FOR_AGENT.md)
- [改进版 API 说明](../README_improved_api.md)