scripts/data/convert_manifeel_to_hf.py

#!/usr/bin/env python3
"""将 ManiFeel (zarr 格式) 转换为 HuggingFace datasets 格式。

用法:
    python scripts/data/convert_manifeel_to_hf.py \\
        --input /path/to/manifeel_zarr \\
        --output /path/to/hf_manifeel \\
        --num-workers 8

ManiFeel zarr 结构:
    <task_dir>/
    ├── front/        # zarr 数组 [T, H, W, C]
    ├── wrist/
    ├── side/
    ├── state/        # [T, 7]
    ├── action/       # [T, 6]
    └── meta/
        └── episode_ends  # episode 边界索引
"""

import os, json, argparse, logging
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import zarr
import cv2
from tqdm import tqdm

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# DreamZero 标准参数
ACTION_HORIZON = 12
MAX_STATE_DIM = 44
MAX_ACTION_DIM = 32
VIDEO_HEIGHT = 160
VIDEO_WIDTH = 320
FPS = 30


def pad_state(state: np.ndarray) -> np.ndarray:
    d = state.shape[-1]
    padded = np.zeros((MAX_STATE_DIM,), dtype=np.float32)
    padded[:d] = state.astype(np.float32)
    return padded


def pad_action(action_chunk: np.ndarray) -> np.ndarray:
    """action_chunk: [horizon, D]"""
    d = action_chunk.shape[-1]
    padded = np.zeros((*action_chunk.shape[:-1], MAX_ACTION_DIM), dtype=np.float32)
    padded[..., :d] = action_chunk.astype(np.float32)
    return padded


def process_task(args):
    """处理一个 ManiFeel 任务目录，返回训练样本列表。"""
    task_dir_str, output_dir = args
    task_dir = Path(task_dir_str)
    task_name = task_dir.name

    try:
        # 读取 zarr
        front = zarr.open(str(task_dir / "front"), mode="r")
        wrist = zarr.open(str(task_dir / "wrist"), mode="r")
        side = zarr.open(str(task_dir / "side"), mode="r")
        state_arr = zarr.open(str(task_dir / "state"), mode="r")
        action_arr = zarr.open(str(task_dir / "action"), mode="r")
        episode_ends = zarr.open(str(task_dir / "meta" / "episode_ends"), mode="r")
    except Exception as e:
        logger.warning(f"无法读取 {task_dir}: {e}")
        return []

    # zarr 读取为 numpy
    front_np = np.array(front)        # [T, H, W, C]
    wrist_np = np.array(wrist)
    side_np = np.array(side)
    state_np = np.array(state_arr)    # [T, D_state]
    action_np = np.array(action_arr)  # [T, D_action]
    ends = np.array(episode_ends)     # episode 边界索引

    records = []
    video_out_dir = output_dir / "videos"
    # 3 views
    for view_idx, view_name in enumerate(["front", "wrist", "side"]):
        (video_out_dir / f"view_{view_idx}").mkdir(parents=True, exist_ok=True)

    prev_end = 0
    for ep_idx, end in enumerate(ends):
        start = prev_end
        prev_end = end
        ep_len = end - start
        if ep_len < ACTION_HORIZON + 1:
            continue

        # 为每个视角编码 mp4
        video_paths = []
        for view_idx, view_data in enumerate([front_np, wrist_np, side_np]):
            ep_frames = view_data[start:end]  # [ep_len, H, W, C]
            video_filename = f"episode_{task_name}_{ep_idx:04d}_view{view_idx}.mp4"
            video_path = video_out_dir / f"view_{view_idx}" / video_filename

            # resize + encode
            h, w = ep_frames.shape[1], ep_frames.shape[2]
            writer = cv2.VideoWriter(
                str(video_path),
                cv2.VideoWriter_fourcc(*"mp4v"),
                FPS, (VIDEO_WIDTH, VIDEO_HEIGHT),
            )
            for f_idx in range(ep_len):
                frame = ep_frames[f_idx]
                if h != VIDEO_HEIGHT or w != VIDEO_WIDTH:
                    frame = cv2.resize(frame, (VIDEO_WIDTH, VIDEO_HEIGHT))
                if frame.shape[-1] == 3:
                    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                writer.write(frame)
            writer.release()
            video_paths.append(f"videos/view_{view_idx}/{video_filename}")

        # 为每个时间窗口生成样本
        ep_state = state_np[start:end]
        ep_action = action_np[start:end]

        for t in range(ep_len - ACTION_HORIZON):
            state = pad_state(ep_state[t])
            action_chunk = ep_action[t:t + ACTION_HORIZON]  # [horizon, D]
            action_chunk = pad_action(action_chunk)

            # ManiFeel 所有视角拼接到一条样本，以 dict 形式存储
            records.append({
                "task": task_name,
                "episode_index": ep_idx,
                "frame_index": t,
                "state": state.tolist(),
                "action": action_chunk.tolist(),
                "action_mask": [True] * ACTION_HORIZON,
                "text": task_name.replace("_", " ").replace("-", " "),
                "video_path_0": video_paths[0],
                "video_path_1": video_paths[1],
                "video_path_2": video_paths[2],
            })

    logger.info(f"  {task_name}: {len(records)} samples")
    return records


def main():
    parser = argparse.ArgumentParser(description="Convert ManiFeel zarr to HF datasets format")
    parser.add_argument("--input", "-i", required=True, help="ManiFeel 数据根目录（包含任务子目录）")
    parser.add_argument("--output", "-o", required=True, help="HF 数据集输出目录")
    parser.add_argument("--num-workers", "-w", type=int, default=4, help="并行任务数")
    args = parser.parse_args()

    input_dir = Path(args.input)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 扫描任务目录
    task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()])
    logger.info(f"找到 {len(task_dirs)} 个任务目录")

    # 处理每个任务
    all_records = []
    for task_dir in tqdm(task_dirs, desc="Processing tasks"):
        records = process_task((str(task_dir), output_dir))
        all_records.extend(records)

    logger.info(f"总共生成 {len(all_records)} 条训练样本")

    if not all_records:
        logger.error("未生成任何样本！")
        return

    # 写入 parquet
    df = pd.DataFrame(all_records)
    data_dir = output_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)

    table = pa.Table.from_pandas(df)
    pq.write_table(table, data_dir / "train-00000.parquet")

    # dataset_info.json
    dataset_info = {
        "description": "ManiFeel benchmark dataset for DreamZero",
        "features": {
            "task": {"dtype": "string", "_type": "Value"},
            "episode_index": {"dtype": "int64", "_type": "Value"},
            "frame_index": {"dtype": "int64", "_type": "Value"},
            "state": {"dtype": "float32", "shape": [MAX_STATE_DIM], "_type": "Sequence"},
            "action": {"dtype": "float32", "shape": [ACTION_HORIZON, MAX_ACTION_DIM], "_type": "Sequence"},
            "action_mask": {"dtype": "bool", "shape": [ACTION_HORIZON], "_type": "Sequence"},
            "text": {"dtype": "string", "_type": "Value"},
            "video_path_0": {"dtype": "string", "_type": "Value"},
            "video_path_1": {"dtype": "string", "_type": "Value"},
            "video_path_2": {"dtype": "string", "_type": "Value"},
        },
        "num_views": 3,
        "splits": {"train": {"num_examples": len(all_records)}},
    }
    with open(output_dir / "dataset_info.json", "w") as f:
        json.dump(dataset_info, f, indent=2)

    # README
    readme = f"""---
license: cc-by-4.0
---

# DreamZero - ManiFeel

## Description
ManiFeel benchmark dataset (3 views: front, wrist, side) converted for DreamZero.

## Schema
| Column | Type | Description |
|--------|------|-------------|
| video_path_0/1/2 | string | Front/wrist/side video files |
| state | float32[{MAX_STATE_DIM}] | Robot state (padded) |
| action | float32[{ACTION_HORIZON},{MAX_ACTION_DIM}] | Action chunks (padded) |
| text | string | Task description |
| episode_index | int64 | Episode ID |

## Statistics
- Total samples: {len(all_records)}
- Views: 3 (front, wrist, side)
- Video resolution: {VIDEO_WIDTH}x{VIDEO_HEIGHT}
"""
    with open(output_dir / "README.md", "w") as f:
        f.write(readme)

    logger.info(f"转换完成！输出: {output_dir}")


if __name__ == "__main__":
    main()