670 lines
28 KiB
Python
670 lines
28 KiB
Python
# source /fs-computility/efm/liyang/miniconda3/etc/profile.d/conda.sh
|
||
# conda activate act
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import gc
|
||
import shutil
|
||
from concurrent.futures import ALL_COMPLETED, ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait
|
||
from pathlib import Path
|
||
from typing import Callable, Dict, List, Optional, Tuple
|
||
import torchvision
|
||
import cv2
|
||
import h5py
|
||
import lmdb
|
||
import numpy as np
|
||
import pickle
|
||
import torch
|
||
from PIL import Image
|
||
from scipy.spatial.transform import Rotation
|
||
from tqdm import tqdm
|
||
import logging
|
||
import pdb
|
||
import os
|
||
import imageio # imageio-ffmpeg
|
||
from lerobot.common.datasets.compute_stats import auto_downsample_height_width, get_feature_stats, sample_indices
|
||
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
||
from lerobot.common.datasets.utils import check_timestamps_sync, get_episode_data_index, validate_episode_buffer
|
||
import time
|
||
# import ray
|
||
# from ray.runtime_env import RuntimeEnv
|
||
|
||
"""
|
||
Store both camera image and robot state as a combined observation.
|
||
Args:
|
||
observation: images(camera), states (robot state)
|
||
actions: joint, gripper, ee_pose
|
||
"""
|
||
FEATURES = {
|
||
"images.rgb.head": {
|
||
"dtype": "video",
|
||
"shape": (368, 640, 3),
|
||
"names": ["height", "width", "channel"],
|
||
},
|
||
"images.rgb.hand_left": {
|
||
"dtype": "video",
|
||
"shape": (480, 640, 3),
|
||
"names": ["height", "width", "channel"],
|
||
},
|
||
"images.rgb.hand_right": {
|
||
"dtype": "video",
|
||
"shape": (480, 640, 3),
|
||
"names": ["height", "width", "channel"],
|
||
},
|
||
# "states.left_joint.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (6,),
|
||
# "names": ["left_joint_0", "left_joint_1", "left_joint_2", "left_joint_3", "left_joint_4", "left_joint_5",],
|
||
# },
|
||
# "states.left_gripper.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (1,),
|
||
# "names": ["left_gripper_0",],
|
||
# },
|
||
# "states.right_joint.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (6,),
|
||
# "names": ["right_joint_0", "right_joint_1", "right_joint_2", "right_joint_3", "right_joint_4", "right_joint_5",],
|
||
# },
|
||
# "states.right_gripper.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (1,),
|
||
# "names": ["right_gripper_0",],
|
||
# },
|
||
"observation.state": {
|
||
"dtype": "float32",
|
||
"shape": (14,),
|
||
"names": ["left_joint_0", "left_joint_1", "left_joint_2", "left_joint_3", "left_joint_4", "left_joint_5", "left_gripper_0",
|
||
"right_joint_0", "right_joint_1", "right_joint_2", "right_joint_3", "right_joint_4", "right_joint_5","right_gripper_0"],
|
||
},
|
||
"action": {
|
||
"dtype": "float32",
|
||
"shape": (14,),
|
||
"names": ["left_joint_0", "left_joint_1", "left_joint_2", "left_joint_3", "left_joint_4", "left_joint_5", "left_gripper_0",
|
||
"right_joint_0", "right_joint_1", "right_joint_2", "right_joint_3", "right_joint_4", "right_joint_5","right_gripper_0"],
|
||
},
|
||
# "actions.left_joint.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (6,),
|
||
# "names": ["left_joint_0", "left_joint_1", "left_joint_2", "left_joint_3", "left_joint_4", "left_joint_5",],
|
||
# },
|
||
# "actions.left_gripper.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (1,),
|
||
# "names": ["left_gripper_0",],
|
||
# },
|
||
# "actions.right_joint.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (6,),
|
||
# "names": ["right_joint_0", "right_joint_1", "right_joint_2", "right_joint_3", "right_joint_4", "right_joint_5",],
|
||
# },
|
||
# "actions.right_gripper.position": {
|
||
# "dtype": "float32",
|
||
# "shape": (1,),
|
||
# "names": ["right_gripper_0", ],
|
||
# },
|
||
|
||
}
|
||
|
||
|
||
import numpy as np
|
||
|
||
def filter_forbidden_frames(state_dict, position_threshold=0.001, velocity_threshold=0.005):
|
||
"""
|
||
过滤禁止的帧，基于位置和速度阈值
|
||
|
||
参数:
|
||
- state_dict: 形状为 (n, 14) 的状态数组
|
||
- position_threshold: 位置变化的阈值
|
||
- velocity_threshold: 速度变化的阈值
|
||
|
||
返回:
|
||
- valid_mask: 布尔数组，True表示有效帧
|
||
"""
|
||
# 排除夹爪列（第6和第13列，索引从0开始）
|
||
qpos_columns = [i for i in range(14)]
|
||
qpos_data = state_dict[:, qpos_columns]
|
||
|
||
n_frames = len(state_dict)
|
||
valid_mask = np.ones(n_frames, dtype=bool)
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
# 计算帧间差异（速度）
|
||
if n_frames > 1:
|
||
|
||
diff_sum = np.sum(np.abs(np.diff(qpos_data, axis=0)), axis=1)
|
||
# sorted_indices = np.argsort(diff_sum)[::-1]
|
||
# sorted_abs_sums = diff_sum[sorted_indices]
|
||
|
||
# velocities = np.diff(qpos_data, axis=0)
|
||
# 检查速度是否超过阈值
|
||
for i in range(n_frames - 1):
|
||
if np.any(np.abs(diff_sum[i]) > position_threshold):
|
||
valid_mask[i] = True # 有运动，有效帧
|
||
else:
|
||
valid_mask[i] = False # 静止，可能是禁止帧
|
||
valid_mask[i] = True
|
||
return valid_mask
|
||
|
||
def statistical_filter(state_dict, std_multiplier=2.0):
|
||
"""
|
||
使用统计方法检测异常（禁止）帧
|
||
"""
|
||
# 排除夹爪列
|
||
qpos_columns = [i for i in range(14) if i not in [6, 13]]
|
||
qpos_data = state_dict[:, qpos_columns]
|
||
|
||
# 计算每列的均值和标准差
|
||
means = np.mean(qpos_data, axis=0)
|
||
stds = np.std(qpos_data, axis=0)
|
||
|
||
# 创建有效掩码
|
||
valid_mask = np.ones(len(state_dict), dtype=bool)
|
||
|
||
for i in range(len(state_dict)):
|
||
# 检查每个关节位置是否在合理范围内
|
||
deviations = np.abs(qpos_data[i] - means)
|
||
if np.any(deviations > std_multiplier * stds):
|
||
valid_mask[i] = False # 异常帧
|
||
|
||
return valid_mask
|
||
|
||
|
||
class ARXDataset(LeRobotDataset):
|
||
def __init__(
|
||
self,
|
||
repo_id: str,
|
||
root: str | Path | None = None,
|
||
episodes: list[int] | None = None,
|
||
image_transforms: Callable | None = None,
|
||
delta_timestamps: dict[list[float]] | None = None,
|
||
tolerance_s: float = 1e-4,
|
||
download_videos: bool = True,
|
||
local_files_only: bool = False,
|
||
video_backend: str | None = None,
|
||
):
|
||
super().__init__(
|
||
repo_id=repo_id,
|
||
root=root,
|
||
episodes=episodes,
|
||
image_transforms=image_transforms,
|
||
delta_timestamps=delta_timestamps,
|
||
tolerance_s=tolerance_s,
|
||
download_videos=download_videos,
|
||
local_files_only=local_files_only,
|
||
video_backend=video_backend,
|
||
)
|
||
|
||
def save_episode(self, episode_data: dict | None = None, videos: dict | None = None) -> None:
|
||
if not episode_data:
|
||
episode_buffer = self.episode_buffer
|
||
|
||
validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
|
||
episode_length = episode_buffer.pop("size")
|
||
tasks = episode_buffer.pop("task")
|
||
episode_tasks = list(set(tasks))
|
||
episode_index = episode_buffer["episode_index"]
|
||
|
||
episode_buffer["index"] = np.arange(self.meta.total_frames, self.meta.total_frames + episode_length)
|
||
episode_buffer["episode_index"] = np.full((episode_length,), episode_index)
|
||
|
||
for task in episode_tasks:
|
||
task_index = self.meta.get_task_index(task)
|
||
if task_index is None:
|
||
self.meta.add_task(task)
|
||
|
||
episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks])
|
||
for key, ft in self.features.items():
|
||
if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["video"]:
|
||
continue
|
||
episode_buffer[key] = np.stack(episode_buffer[key]).squeeze()
|
||
for key in self.meta.video_keys:
|
||
video_path = self.root / self.meta.get_video_file_path(episode_index, key)
|
||
episode_buffer[key] = str(video_path) # PosixPath -> str
|
||
video_path.parent.mkdir(parents=True, exist_ok=True)
|
||
shutil.copyfile(videos[key], video_path)
|
||
ep_stats = compute_episode_stats(episode_buffer, self.features)
|
||
self._save_episode_table(episode_buffer, episode_index)
|
||
self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
|
||
ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index])
|
||
ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()}
|
||
check_timestamps_sync(
|
||
episode_buffer["timestamp"],
|
||
episode_buffer["episode_index"],
|
||
ep_data_index_np,
|
||
self.fps,
|
||
self.tolerance_s,
|
||
)
|
||
if not episode_data:
|
||
self.episode_buffer = self.create_episode_buffer()
|
||
|
||
|
||
def add_frame(self, frame: dict) -> None:
|
||
for name in frame:
|
||
if isinstance(frame[name], torch.Tensor):
|
||
frame[name] = frame[name].numpy()
|
||
features = {key: value for key, value in self.features.items() if key in self.hf_features}
|
||
if self.episode_buffer is None:
|
||
self.episode_buffer = self.create_episode_buffer()
|
||
frame_index = self.episode_buffer["size"]
|
||
timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
|
||
self.episode_buffer["frame_index"].append(frame_index)
|
||
self.episode_buffer["timestamp"].append(timestamp)
|
||
|
||
for key in frame:
|
||
if key == "task":
|
||
self.episode_buffer["task"].append(frame["task"])
|
||
continue
|
||
if key not in self.features:
|
||
print("key ", key)
|
||
raise ValueError(f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'.")
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
self.episode_buffer[key].append(frame[key])
|
||
self.episode_buffer["size"] += 1
|
||
|
||
# def crop_resize_no_padding(image, target_size=(480, 640)):
|
||
# """
|
||
# Crop and scale to target size (no padding)
|
||
# :param image: input image (NumPy array)
|
||
# :param target_size: target size (height, width)
|
||
# :return: processed image
|
||
# """
|
||
# h, w = image.shape[:2]
|
||
# target_h, target_w = target_size
|
||
# target_ratio = target_w / target_h # Target aspect ratio (e.g. 640/480=1.333)
|
||
|
||
# # the original image aspect ratio and cropping direction
|
||
# if w / h > target_ratio: # Original image is wider → crop width
|
||
# crop_w = int(h * target_ratio) # Calculate crop width based on target aspect ratio
|
||
# crop_h = h
|
||
# start_x = (w - crop_w) // 2 # Horizontal center starting point
|
||
# start_y = 0
|
||
# else: # Original image is higher → crop height
|
||
# crop_h = int(w / target_ratio) # Calculate clipping height according to target aspect ratio
|
||
# crop_w = w
|
||
# start_x = 0
|
||
# start_y = (h - crop_h) // 2 # Vertical center starting point
|
||
|
||
# # Perform centered cropping (to prevent out-of-bounds)
|
||
# start_x, start_y = max(0, start_x), max(0, start_y)
|
||
# end_x, end_y = min(w, start_x + crop_w), min(h, start_y + crop_h)
|
||
# cropped = image[start_y:end_y, start_x:end_x]
|
||
|
||
# # Resize to target size (bilinear interpolation)
|
||
# resized = cv2.resize(cropped, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
|
||
# return resized
|
||
|
||
|
||
def load_lmdb_data(episode_path: Path, sava_path: Path, fps_factor: int, target_fps: int) -> Optional[Dict]:
|
||
def load_image(txn, key):
|
||
raw = txn.get(key)
|
||
data = pickle.loads(raw)
|
||
image = cv2.imdecode(data, cv2.IMREAD_COLOR)
|
||
# Convert to RGB if necessary
|
||
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
# image = crop_resize_no_padding(image, target_size=(480, 640))
|
||
return image
|
||
try:
|
||
env = lmdb.open(
|
||
str(episode_path / "lmdb"),
|
||
readonly=True,
|
||
lock=False,
|
||
max_readers=128,
|
||
readahead=False
|
||
)
|
||
with env.begin(write=False) as txn:
|
||
keys = [k for k, _ in txn.cursor()]
|
||
|
||
image_keys = sorted([k for k in keys if b'head' in k])
|
||
if not image_keys:
|
||
return None
|
||
|
||
all_qpos = pickle.loads(txn.get(b'/observations/qpos'))
|
||
|
||
if np.isscalar(all_qpos):
|
||
total_steps = len(image_keys)
|
||
all_qpos = [all_qpos] * total_steps
|
||
else:
|
||
total_steps = len(all_qpos)
|
||
all_qpos = np.stack(all_qpos)
|
||
state_action_dict = {}
|
||
state_action_dict["states.left_joint.position"] = all_qpos[:, :6]
|
||
state_action_dict["states.left_gripper.position"] = all_qpos[:, 6][:, None] # np.expand_dims(all_qpos[:, 6], axis=1)
|
||
state_action_dict["states.right_joint.position"] = all_qpos[:, 7:13]
|
||
state_action_dict["states.right_gripper.position"] = all_qpos[:, 13][:, None] # np.expand_dims(all_qpos[:, 13], axis=1)
|
||
# state_keys = list(state_action_dict.keys())
|
||
# for k in state_keys:
|
||
# state_action_dict[k.replace("states", "actions")] = np.concatenate([state_action_dict[k][1:, :], state_action_dict[k][-1, :][None,:]], axis=0)
|
||
|
||
|
||
# action_dict = {}
|
||
# action_dict["actions.left_joint.position"] = np.concatenate([state_dict["states.left_joint.position"][1:, :], state_dict["states.left_joint.position"][-1, :][None,:]], axis=0)
|
||
# action_dict["actions.left_gripper.position"] = state_dict["states.left_gripper.position"][1:, :]
|
||
# action_dict["actions.right_joint.position"] = state_dict["states.right_joint.position"][1:, :]
|
||
# action_dict["actions.right_gripper.position"] = state_dict["states.right_gripper.position"][1:, :]
|
||
|
||
action_dict = {}
|
||
|
||
action_dict["action"] = np.concatenate([all_qpos[1:,], all_qpos[-1,].reshape(-1, 14)], axis=0)
|
||
state_dict = {}
|
||
state_dict["observation.state"] = all_qpos
|
||
mask1 = filter_forbidden_frames(state_dict["observation.state"])
|
||
# state_dict["observation.state"] = state_dict["observation.state"][mask1]
|
||
# action_dict["actions.left_gripper.position"] = state_dict["states.left_gripper.position"][1:, :]
|
||
# action_dict["actions.right_arm.position"] = np.concatenate([state_action_dict["states.right_joint.position"][1:, :], state_action_dict["states.right_joint.position"][-1, :][None,:]], axis=0)
|
||
# action_dict["actions.left_arm.position"] = state_dict["states.right_gripper.position"][1:, :]
|
||
|
||
assert total_steps == len(image_keys), "qpos length mismatch"
|
||
selected_steps = [step for step in range(total_steps) if step % fps_factor == 0 and mask1[step]]
|
||
frames = []
|
||
image_observations = {
|
||
"images.rgb.head": [],
|
||
"images.rgb.hand_left": [],
|
||
"images.rgb.hand_right": []
|
||
}
|
||
|
||
start_time = time.time()
|
||
|
||
for step_index, step in enumerate(selected_steps):
|
||
step_str = f"{step:04d}"
|
||
head_key = f"observation/head/color_image/{step_str}".encode()
|
||
left_key = f"observation/left_wrist/color_image/{step_str}".encode()
|
||
right_key = f"observation/right_wrist/color_image/{step_str}".encode()
|
||
if not (head_key in keys and left_key in keys and right_key in keys):
|
||
continue
|
||
# state = all_qpos[step]
|
||
# if step_index < len(selected_steps) - 1:
|
||
# action = all_qpos[selected_steps[step_index + 1]]
|
||
# else:
|
||
# action = state
|
||
data_dict = {}
|
||
# for key, value in state_action_dict.items():
|
||
# data_dict[key] = value[step]
|
||
data_dict['action'] = action_dict["action"][step]
|
||
data_dict["task"] = " ".join(episode_path.parent.parent.name.split("_"))
|
||
data_dict['observation.state'] = state_dict["observation.state"][step]
|
||
# frames.append({
|
||
# "observation.states.joint.position": state,
|
||
# "actions.joint.position": action,
|
||
# "task": task_name,
|
||
# })
|
||
frames.append(data_dict)
|
||
image_observations["images.rgb.head"].append(load_image(txn, head_key))
|
||
image_observations["images.rgb.hand_left"].append(load_image(txn, left_key))
|
||
image_observations["images.rgb.hand_right"].append(load_image(txn, right_key))
|
||
end_time = time.time()
|
||
elapsed_time = end_time - start_time
|
||
print(f"load image_observations of {episode_path}")
|
||
env.close()
|
||
if not frames:
|
||
return None
|
||
os.makedirs(sava_path, exist_ok=True)
|
||
os.makedirs(sava_path/episode_path.name, exist_ok=True)
|
||
imageio.mimsave(sava_path/episode_path.name/'head.mp4', image_observations["images.rgb.head"], fps=target_fps)
|
||
imageio.mimsave(sava_path/episode_path.name/'hand_left.mp4', image_observations["images.rgb.hand_left"], fps=target_fps)
|
||
imageio.mimsave(sava_path/episode_path.name/'hand_right.mp4', image_observations["images.rgb.hand_right"], fps=target_fps)
|
||
print(f"imageio.mimsave time taken of {episode_path}")
|
||
|
||
return {
|
||
"frames": frames,
|
||
"videos": {
|
||
"images.rgb.head": sava_path/episode_path.name/"head.mp4",
|
||
"images.rgb.hand_left": sava_path/episode_path.name/"hand_left.mp4",
|
||
"images.rgb.hand_right": sava_path/episode_path.name/"hand_right.mp4",
|
||
},
|
||
}
|
||
|
||
except Exception as e:
|
||
logging.error(f"Failed to load LMDB data: {e}")
|
||
return None
|
||
|
||
|
||
def get_all_tasks(src_path: Path, output_path: Path) -> Tuple[Path, Path]:
|
||
src_dirs = sorted(list(src_path.glob("*"))) # "set*-*_collector*_datatime" as the conversion unit
|
||
|
||
save_dirs = [output_path/_dir.parent.name/_dir.name for _dir in src_dirs]
|
||
tasks_tuples = zip(src_dirs, save_dirs)
|
||
for task in tasks_tuples:
|
||
yield task
|
||
|
||
def compute_episode_stats(episode_data: Dict[str, List[str] | np.ndarray], features: Dict) -> Dict:
|
||
ep_stats = {}
|
||
for key, data in episode_data.items():
|
||
if features[key]["dtype"] == "string":
|
||
continue
|
||
elif features[key]["dtype"] in ["image", "video"]:
|
||
ep_ft_array = sample_images(data)
|
||
axes_to_reduce = (0, 2, 3) # keep channel dim
|
||
keepdims = True
|
||
else:
|
||
ep_ft_array = data # data is already a np.ndarray
|
||
axes_to_reduce = 0 # compute stats over the first axis
|
||
keepdims = data.ndim == 1 # keep as np.array
|
||
|
||
ep_stats[key] = get_feature_stats(ep_ft_array, axis=axes_to_reduce, keepdims=keepdims)
|
||
if features[key]["dtype"] in ["image", "video"]:
|
||
ep_stats[key] = {
|
||
k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items()
|
||
}
|
||
return ep_stats
|
||
|
||
def sample_images(input):
|
||
if type(input) is str:
|
||
video_path = input
|
||
reader = torchvision.io.VideoReader(video_path, stream="video")
|
||
frames = [frame["data"] for frame in reader]
|
||
frames_array = torch.stack(frames).numpy() # Shape: [T, C, H, W]
|
||
sampled_indices = sample_indices(len(frames_array))
|
||
images = None
|
||
for i, idx in enumerate(sampled_indices):
|
||
img = frames_array[idx]
|
||
img = auto_downsample_height_width(img)
|
||
if images is None:
|
||
images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
|
||
images[i] = img
|
||
elif type(input) is np.ndarray:
|
||
frames_array = input[:, None, :, :] # Shape: [T, C, H, W]
|
||
sampled_indices = sample_indices(len(frames_array))
|
||
images = None
|
||
for i, idx in enumerate(sampled_indices):
|
||
img = frames_array[idx]
|
||
img = auto_downsample_height_width(img)
|
||
if images is None:
|
||
images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
|
||
images[i] = img
|
||
return images
|
||
|
||
|
||
def load_local_dataset(episode_path: str, save_path:str, origin_fps=30, target_fps=30):
|
||
fps_factor = origin_fps // target_fps
|
||
# print(f"fps downsample factor: {fps_factor}")
|
||
# logging.info(f"fps downsample factor: {fps_factor}")
|
||
# for format_str in [f"{episode_id:07d}", f"{episode_id:06d}", str(episode_id)]:
|
||
# episode_path = Path(src_path) / format_str
|
||
# save_path = Path(save_path) / format_str
|
||
# if episode_path.exists():
|
||
# break
|
||
# else:
|
||
# logging.warning(f"Episode directory not found for ID {episode_id}")
|
||
# return None, None
|
||
episode_path = Path(episode_path)
|
||
if not episode_path.exists():
|
||
logging.warning(f"{episode_path} does not exist")
|
||
return None, None
|
||
|
||
if not (episode_path / "lmdb/data.mdb").exists():
|
||
logging.warning(f"LMDB data not found for episode {episode_path}")
|
||
return None, None
|
||
|
||
raw_dataset = load_lmdb_data(episode_path, save_path, fps_factor, target_fps)
|
||
if raw_dataset is None:
|
||
return None, None
|
||
frames = raw_dataset["frames"] # states, actions, task
|
||
|
||
videos = raw_dataset["videos"] # image paths
|
||
## check the frames
|
||
for camera_name, video_path in videos.items():
|
||
if not os.path.exists(video_path):
|
||
logging.error(f"Video file {video_path} does not exist.")
|
||
print(f"Camera {camera_name} Video file {video_path} does not exist.")
|
||
return None, None
|
||
return frames, videos
|
||
|
||
|
||
def save_as_lerobot_dataset(task: tuple[Path, Path], repo_id, num_threads, debug, origin_fps=30, target_fps=30, robot_type="piper", delete_downsampled_videos=True):
|
||
src_path, save_path = task
|
||
print(f"**Processing collected** {src_path}")
|
||
print(f"**saving to** {save_path}")
|
||
if save_path.exists():
|
||
# print(f"Output directory {save_path} already exists. Deleting it.")
|
||
# logging.warning(f"Output directory {save_path} already exists. Deleting it.")
|
||
# shutil.rmtree(save_path)
|
||
print(f"Output directory {save_path} already exists.")
|
||
return
|
||
|
||
dataset = ARXDataset.create(
|
||
repo_id=f"{repo_id}",
|
||
root=save_path,
|
||
fps=target_fps,
|
||
robot_type=robot_type,
|
||
features=FEATURES,
|
||
)
|
||
all_episode_paths = sorted([f.as_posix() for f in src_path.glob(f"*") if f.is_dir()])
|
||
# all_subdir_eids = [int(Path(path).name) for path in all_subdir]
|
||
if debug:
|
||
for i in range(1):
|
||
# pdb.set_trace()
|
||
frames, videos = load_local_dataset(episode_path=all_episode_paths[i], save_path=save_path, origin_fps=origin_fps, target_fps=target_fps)
|
||
for frame_data in frames:
|
||
dataset.add_frame(frame_data)
|
||
dataset.save_episode(videos=videos)
|
||
if delete_downsampled_videos:
|
||
for _, video_path in videos.items():
|
||
parent_dir = os.path.dirname(video_path)
|
||
try:
|
||
shutil.rmtree(parent_dir)
|
||
# os.remove(video_path)
|
||
# print(f"Successfully deleted: {parent_dir}")
|
||
print(f"Successfully deleted: {video_path}")
|
||
except Exception as e:
|
||
pass # Handle the case where the directory might not exist or is already deleted
|
||
else:
|
||
for batch_index in range(len(all_episode_paths)//num_threads+1):
|
||
batch_episode_paths = all_episode_paths[batch_index*num_threads:(batch_index+1)*num_threads]
|
||
if len(batch_episode_paths) == 0:
|
||
continue
|
||
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
||
futures = []
|
||
for episode_path in batch_episode_paths:
|
||
print("starting to process episode: ", episode_path)
|
||
futures.append(
|
||
executor.submit(load_local_dataset, episode_path=episode_path, save_path=save_path, origin_fps=origin_fps, target_fps=target_fps)
|
||
)
|
||
for raw_dataset in as_completed(futures):
|
||
frames, videos = raw_dataset.result()
|
||
if frames is None or videos is None:
|
||
print(f"Skipping episode {episode_path} due to missing data.")
|
||
continue
|
||
for frame_data in frames:
|
||
dataset.add_frame(frame_data)
|
||
dataset.save_episode(videos=videos)
|
||
gc.collect()
|
||
print(f"finishing processed {videos}")
|
||
if delete_downsampled_videos:
|
||
for _, video_path in videos.items():
|
||
# Get the parent directory of the video
|
||
parent_dir = os.path.dirname(video_path)
|
||
try:
|
||
shutil.rmtree(parent_dir)
|
||
print(f"Successfully deleted: {parent_dir}")
|
||
except Exception as e:
|
||
pass
|
||
|
||
def main(src_path, save_path, repo_id, num_threads=60, debug=False, origin_fps=30, target_fps=30):
|
||
logging.info("Scanning for episodes...")
|
||
tasks = get_all_tasks(src_path, save_path)
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
if debug:
|
||
task = next(tasks)
|
||
save_as_lerobot_dataset(task, repo_id, num_threads=num_threads, debug=debug, origin_fps=origin_fps, target_fps=target_fps)
|
||
else:
|
||
for task in tasks:
|
||
save_as_lerobot_dataset(task, repo_id, num_threads=num_threads, debug=debug, origin_fps=origin_fps, target_fps=target_fps)
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="Convert collected data from Piper to Lerobot format.")
|
||
parser.add_argument(
|
||
"--src_path",
|
||
type=str,
|
||
# required=False,
|
||
default="/fs-computility/efm/shared/datasets/myData-A1/real/raw_data/agilex_split_aloha/",
|
||
help="Path to the input file containing collected data in Piper format.",
|
||
#help="/fs-computility/efm/shared/datasets/myData-A1/real/raw_data/agilex_split_aloha/Make_a_beef_sandwich",
|
||
)
|
||
parser.add_argument(
|
||
"--save_path",
|
||
type=str,
|
||
# required=False,
|
||
default="/fs-computility/efm/shared/datasets/myData-A1/real/lerobot_v2_1/agilex_split_aloha/",
|
||
help="Path to the output file where the converted Lerobot format will be saved.",
|
||
#help="Path to the output file where the converted Lerobot format will be saved.",
|
||
)
|
||
parser.add_argument(
|
||
"--debug",
|
||
action="store_true",
|
||
help="Run in debug mode with limited episodes",
|
||
)
|
||
parser.add_argument(
|
||
"--num-threads",
|
||
type=int,
|
||
default=50,
|
||
help="Number of threads per process",
|
||
)
|
||
# parser.add_argument(
|
||
# "--task_name",
|
||
# type=str,
|
||
# required=True,
|
||
# default="Pick_up_the_marker_and_put_it_into_the_pen_holder",
|
||
# help="Name of the task to be processed. Default is 'Pick_up_the_marker_and_put_it_into_the_pen_holder'.",
|
||
# )
|
||
parser.add_argument(
|
||
"--repo_id",
|
||
type=str,
|
||
required=True,
|
||
# default="SplitAloha_20250714",
|
||
help="identifier for the dataset repository.",
|
||
)
|
||
parser.add_argument(
|
||
"--origin_fps",
|
||
type=int,
|
||
default=30,
|
||
help="Frames per second for the obervation video. Default is 30.",
|
||
)
|
||
parser.add_argument(
|
||
"--target_fps",
|
||
type=int,
|
||
default=30,
|
||
help="Frames per second for the downsample video. Default is 30.",
|
||
)
|
||
args = parser.parse_args()
|
||
assert int(args.origin_fps) % int(args.target_fps) == 0, "origin_fps must be an integer multiple of target_fps"
|
||
start_time = time.time()
|
||
main(
|
||
src_path=Path(args.src_path),
|
||
save_path=Path(args.save_path),
|
||
repo_id=args.repo_id,
|
||
num_threads=args.num_threads,
|
||
debug=args.debug,
|
||
origin_fps=args.origin_fps,
|
||
target_fps=args.target_fps
|
||
)
|
||
end_time = time.time()
|
||
elapsed_time = end_time - start_time
|
||
print(f"Total time taken: {elapsed_time:.2f} seconds")
|
||
# --target_fps 10
|
||
# --src_path /fs-computility/efm/shared/datasets/myData-A1/real/raw_data/agilex_split_aloha/Put_the_bananas_in_the_basket
|
||
# --save_path /mnt/shared-storage-user/internvla/Users/liyang/data/processed_data/arx_lift2 |