369 lines
14 KiB
Python
369 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
LeRobot到RDT数据转换脚本
|
||
|
||
LeRobot机器人结构：
|
||
- 5个关节 (shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll)
|
||
- 1个夹爪 (gripper)
|
||
- 总计：6个自由度 (6DOF)
|
||
|
||
维度映射（匹配RDT训练代码）：
|
||
- left_arm_dim = 0 (单臂机器人，左臂不存在)
|
||
- right_arm_dim = 6 (5关节 + 1夹爪，映射到RDT的right_arm部分)
|
||
- 状态向量：6维 [joint1, joint2, joint3, joint4, joint5, gripper]
|
||
- RDT索引映射：right_arm_joint_0_pos到right_arm_joint_5_pos (索引0-5)
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import h5py
|
||
import numpy as np
|
||
import cv2
|
||
import argparse
|
||
import yaml
|
||
import json
|
||
import subprocess
|
||
from pathlib import Path
|
||
import pandas as pd
|
||
import torch
|
||
|
||
current_dir = os.path.dirname(__file__)
|
||
sys.path.append(os.path.join(current_dir, ".."))
|
||
from models.multimodal_encoder.t5_encoder import T5Embedder
|
||
|
||
def extract_frames_from_video(video_path, output_dir, episode_idx):
|
||
if not os.path.exists(video_path):
|
||
print(f" No video file: {video_path}")
|
||
return []
|
||
|
||
temp_dir = os.path.join(output_dir, f"temp_frames_{episode_idx}")
|
||
if not os.path.exists(temp_dir):
|
||
os.makedirs(temp_dir)
|
||
|
||
output_pattern = os.path.join(temp_dir, "frame_%04d.jpg")
|
||
|
||
try:
|
||
cmd = [
|
||
'ffmpeg', '-i', video_path,
|
||
'-vf', 'fps=30',
|
||
'-q:v', '2',
|
||
output_pattern,
|
||
'-y'
|
||
]
|
||
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
|
||
if result.returncode != 0:
|
||
print(f" Failed to extract frames with ffmpeg: {result.stderr}")
|
||
return []
|
||
|
||
frames = []
|
||
frame_files = sorted([f for f in os.listdir(temp_dir) if f.endswith('.jpg')])
|
||
|
||
for frame_file in frame_files:
|
||
frame_path = os.path.join(temp_dir, frame_file)
|
||
frame = cv2.imread(frame_path)
|
||
if frame is not None:
|
||
frame_resized = cv2.resize(frame, (640, 480))
|
||
frames.append(frame_resized)
|
||
|
||
print(f" Successfully extracted {len(frames)} frames")
|
||
|
||
for frame_file in frame_files:
|
||
os.remove(os.path.join(temp_dir, frame_file))
|
||
os.rmdir(temp_dir)
|
||
|
||
return frames
|
||
|
||
except Exception as e:
|
||
print(f" Error extracting frames: {e}")
|
||
return []
|
||
|
||
def load_lerobot_episode(data_dir, episode_idx, output_dir, cam_high_key="high", cam_right_wrist_key="arm"):
|
||
"""加载LeRobot的单个episode数据
|
||
|
||
LeRobot数据结构：
|
||
- action: 6维 [shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper]
|
||
- observation.state: 6维 [shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper]
|
||
- 图像: 高位相机 + 手臂相机
|
||
"""
|
||
parquet_path = os.path.join(data_dir, "data/chunk-000", f"episode_{episode_idx:06d}.parquet")
|
||
if not os.path.exists(parquet_path):
|
||
print(f"Episode {episode_idx} parquet file does not exist: {parquet_path}")
|
||
return None
|
||
|
||
df = pd.read_parquet(parquet_path)
|
||
|
||
actions = []
|
||
qpos = []
|
||
|
||
for i in range(len(df)):
|
||
action = df['action'].iloc[i]
|
||
state = df['observation.state'].iloc[i]
|
||
|
||
if isinstance(action, np.ndarray):
|
||
actions.append(action.astype(np.float32))
|
||
else:
|
||
actions.append(np.array(action, dtype=np.float32))
|
||
|
||
if isinstance(state, np.ndarray):
|
||
qpos.append(state.astype(np.float32))
|
||
else:
|
||
qpos.append(np.array(state, dtype=np.float32))
|
||
|
||
high_cam_path = os.path.join(data_dir, f"videos/chunk-000/observation.images.{cam_high_key}", f"episode_{episode_idx:06d}.mp4")
|
||
arm_cam_path = os.path.join(data_dir, f"videos/chunk-000/observation.images.{cam_right_wrist_key}", f"episode_{episode_idx:06d}.mp4")
|
||
|
||
print(f" Extracting high camera frames...")
|
||
high_images = extract_frames_from_video(high_cam_path, output_dir, episode_idx)
|
||
|
||
print(f" Extracting arm camera frames...")
|
||
arm_images = extract_frames_from_video(arm_cam_path, output_dir, episode_idx)
|
||
|
||
target_frames = len(df)
|
||
if len(high_images) > target_frames:
|
||
high_images = high_images[:target_frames]
|
||
if len(arm_images) > target_frames:
|
||
arm_images = arm_images[:target_frames]
|
||
|
||
while len(high_images) < target_frames and high_images:
|
||
high_images.append(high_images[-1])
|
||
while len(arm_images) < target_frames and arm_images:
|
||
arm_images.append(arm_images[-1])
|
||
|
||
return {
|
||
'actions': np.array(actions),
|
||
'qpos': np.array(qpos),
|
||
'high_images': high_images,
|
||
'arm_images': arm_images,
|
||
'episode_length': len(df)
|
||
}
|
||
|
||
def images_encoding(imgs):
|
||
if not imgs:
|
||
return [], 0
|
||
|
||
encode_data = []
|
||
padded_data = []
|
||
max_len = 0
|
||
|
||
for i in range(len(imgs)):
|
||
success, encoded_image = cv2.imencode(".jpg", imgs[i])
|
||
if success:
|
||
jpeg_data = encoded_image.tobytes()
|
||
encode_data.append(jpeg_data)
|
||
max_len = max(max_len, len(jpeg_data))
|
||
else:
|
||
print(f" Image encoding failed: {i}")
|
||
empty_data = b""
|
||
encode_data.append(empty_data)
|
||
|
||
for i in range(len(imgs)):
|
||
padded_data.append(encode_data[i].ljust(max_len, b"\0"))
|
||
|
||
return encode_data, max_len
|
||
|
||
def load_task_instructions(data_dir):
|
||
tasks_file = os.path.join(data_dir, "meta/tasks.jsonl")
|
||
if not os.path.exists(tasks_file):
|
||
print(f"Warning: tasks file not found: {tasks_file}")
|
||
return None
|
||
|
||
instructions = []
|
||
with open(tasks_file, 'r') as f:
|
||
for line in f:
|
||
if line.strip():
|
||
task_data = json.loads(line.strip())
|
||
instructions.append(task_data["task"])
|
||
|
||
print(f" 加载了 {len(instructions)} 个任务指令")
|
||
return instructions
|
||
|
||
def encode_language_instruction(instruction_text, t5_embedder, device):
|
||
try:
|
||
text_embeds, attn_mask = t5_embedder.get_text_embeddings([instruction_text])
|
||
|
||
valid_embeds = text_embeds[0][attn_mask[0]].float()
|
||
return valid_embeds.cpu().numpy()
|
||
|
||
except Exception as e:
|
||
print(f" Language encoding failed: {e}")
|
||
return np.zeros((1, 4096))
|
||
|
||
def convert_lerobot_to_rdt(data_dir, output_dir, episode_num, gpu=0, no_language=False, t5_path=None, cam_high_key="high", cam_right_wrist_key="arm"):
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
print(f"Start converting LeRobot data to RDT format...")
|
||
print(f"Data source: {data_dir}")
|
||
print(f"Output directory: {output_dir}")
|
||
print(f"Processing episode number: {episode_num}")
|
||
print(f"GPU device: {gpu}")
|
||
|
||
scene_name = os.path.basename(data_dir)
|
||
|
||
instructions = None
|
||
if not no_language:
|
||
instructions = load_task_instructions(data_dir)
|
||
|
||
t5_embedder = None
|
||
if not no_language and instructions:
|
||
try:
|
||
print(f" Initializing T5 encoder...")
|
||
t5_embedder = T5Embedder(
|
||
from_pretrained=t5_path,
|
||
device=f"cuda:{gpu}" if torch.cuda.is_available() else "cpu",
|
||
model_max_length=1024,
|
||
use_offload_folder=None,
|
||
)
|
||
print(f" T5 encoder initialized successfully")
|
||
except Exception as e:
|
||
print(f" T5 encoder initialization failed: {e}")
|
||
print(f" Will skip language processing")
|
||
no_language = True
|
||
|
||
for i in range(episode_num):
|
||
print(f"Processing episode {i}...")
|
||
|
||
episode_data = load_lerobot_episode(data_dir, i, output_dir, cam_high_key=cam_high_key, cam_right_wrist_key=cam_right_wrist_key)
|
||
if episode_data is None:
|
||
print(f"Skipping episode {i}")
|
||
continue
|
||
|
||
episode_output_dir = os.path.join(output_dir, f"episode_{i}")
|
||
if not os.path.exists(episode_output_dir):
|
||
os.makedirs(episode_output_dir)
|
||
|
||
hdf5_path = os.path.join(episode_output_dir, f"episode_{i}.hdf5")
|
||
|
||
with h5py.File(hdf5_path, "w") as f:
|
||
f.create_dataset("action", data=episode_data['actions'])
|
||
|
||
obs = f.create_group("observations")
|
||
obs.create_dataset("qpos", data=episode_data['qpos'])
|
||
|
||
image = obs.create_group("images")
|
||
|
||
if episode_data['high_images']:
|
||
print(f" Encoding high camera images...")
|
||
high_enc, len_high = images_encoding(episode_data['high_images'])
|
||
if high_enc and len_high > 0:
|
||
image.create_dataset("cam_high", data=high_enc, dtype=f"S{len_high}")
|
||
print(f" Saved high camera images: {len(episode_data['high_images'])} frames")
|
||
else:
|
||
print(f" Warning: High camera images encoding failed")
|
||
|
||
if episode_data['arm_images']:
|
||
print(f" Encoding arm camera images...")
|
||
arm_enc, len_arm = images_encoding(episode_data['arm_images'])
|
||
if arm_enc and len_arm > 0:
|
||
image.create_dataset("cam_right_wrist", data=arm_enc, dtype=f"S{len_arm}")
|
||
print(f" Saved arm camera images: {len(episode_data['arm_images'])} frames")
|
||
else:
|
||
print(f" Warning: Arm camera images encoding failed")
|
||
|
||
# 添加机器人维度信息（LeRobot: 5个关节 + 1个夹爪）
|
||
# 根据process_data.py的逻辑，每个时间步都需要记录维度信息
|
||
# LeRobot是单臂机器人，只有右臂：5个关节 + 1个夹爪 = 6维
|
||
# 左臂：0维（单臂机器人）
|
||
|
||
# 为每个时间步记录维度信息
|
||
left_arm_dim = [0] * len(episode_data['actions']) # 左臂0维（单臂机器人）
|
||
right_arm_dim = [6] * len(episode_data['actions']) # 右臂6维（5关节+1夹爪）
|
||
|
||
obs.create_dataset("left_arm_dim", data=np.array(left_arm_dim))
|
||
obs.create_dataset("right_arm_dim", data=np.array(right_arm_dim))
|
||
|
||
print(f" Episode {i} converted successfully: {hdf5_path}")
|
||
print(f" Data length: {episode_data['episode_length']}")
|
||
print(f" Action shape: {episode_data['actions'].shape}")
|
||
print(f" Qpos shape: {episode_data['qpos'].shape}")
|
||
print(f" High camera frames: {len(episode_data['high_images'])}")
|
||
print(f" Arm camera frames: {len(episode_data['arm_images'])}")
|
||
|
||
if not no_language and t5_embedder and instructions:
|
||
print(f" Processing language instructions...")
|
||
try:
|
||
instruction = instructions[0]
|
||
|
||
language_features = encode_language_instruction(instruction, t5_embedder, f"cuda:{gpu}")
|
||
|
||
instructions_dir = os.path.join(episode_output_dir, "instructions")
|
||
if not os.path.exists(instructions_dir):
|
||
os.makedirs(instructions_dir)
|
||
|
||
lang_embed_path = os.path.join(instructions_dir, "lang_embed_0.pt")
|
||
torch.save(torch.from_numpy(language_features), lang_embed_path)
|
||
|
||
print(f" Language instruction encoded successfully: {instruction}")
|
||
print(f" Language features saved to: {lang_embed_path}")
|
||
print(f" Language features shape: {language_features.shape}, data type: {language_features.dtype}")
|
||
|
||
except Exception as e:
|
||
print(f" Language instruction processing failed: {e}")
|
||
|
||
print(f"\nConversion completed! Processed {episode_num} episodes")
|
||
print(f"Output directory: {output_dir}")
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Convert LeRobot data to RDT format")
|
||
parser.add_argument("--data_dir", type=str, required=True,
|
||
help="LeRobot data directory path")
|
||
parser.add_argument("--output_dir", type=str, required=True,
|
||
help="Output directory path")
|
||
parser.add_argument("--episode_num", type=int, default=10,
|
||
help="Number of episodes to process")
|
||
parser.add_argument("--gpu", type=int, default=0,
|
||
help="GPU device ID")
|
||
parser.add_argument("--no_language", action="store_true",
|
||
help="Skip language processing")
|
||
parser.add_argument("--cam_high_key", type=str, default="cam_high",
|
||
help="High camera key")
|
||
parser.add_argument("--cam_right_wrist_key", type=str, default="cam_right_wrist",
|
||
help="Right wrist camera key")
|
||
parser.add_argument("--cam_left_wrist_key", type=str, default="cam_left_wrist",
|
||
help="Left wrist camera key")
|
||
parser.add_argument("--t5_path", type=str, required=True,
|
||
help="T5 model path")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if not os.path.exists(args.data_dir):
|
||
print(f"Error: Data directory does not exist: {args.data_dir}")
|
||
return
|
||
|
||
meta_file = os.path.join(args.data_dir, "meta/info.json")
|
||
if not os.path.exists(meta_file):
|
||
print(f"Error: Meta information file not found: {meta_file}")
|
||
return
|
||
|
||
try:
|
||
subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
|
||
print("ffmpeg is available, will use ffmpeg to extract video frames")
|
||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||
print("Warning: ffmpeg is not available, image data may not be extracted correctly")
|
||
print("Please install ffmpeg: conda install -c conda-forge ffmpeg=6.1")
|
||
return
|
||
|
||
with open(meta_file, 'r') as f:
|
||
meta_info = yaml.safe_load(f)
|
||
|
||
total_episodes = meta_info.get('total_episodes', 10)
|
||
if args.episode_num > total_episodes:
|
||
print(f"Warning: Requested episode number ({args.episode_num}) exceeds available number ({total_episodes})")
|
||
args.episode_num = total_episodes
|
||
|
||
convert_lerobot_to_rdt(
|
||
args.data_dir,
|
||
args.output_dir,
|
||
args.episode_num,
|
||
args.gpu,
|
||
args.no_language,
|
||
args.t5_path,
|
||
args.cam_high_key,
|
||
args.cam_right_wrist_key,
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|