#!/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 结构: / ├── 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()