| |
| |
| import os, re |
| import time |
| import numpy as np |
| import h5py |
| import argparse |
| import dm_env |
| import collections |
| from collections import deque |
| import rospy |
| from sensor_msgs.msg import JointState, Image |
| from nav_msgs.msg import Odometry |
| from cv_bridge import CvBridge |
| import cv2 |
| import threading |
| import av |
| import logging |
| from pathlib import Path |
|
|
| def encode_video_frames(images: np.ndarray, dst: Path, fps: int, vcodec: str = "libx264", |
| pix_fmt: str = "yuv420p", g: int = 2, crf: int = 23, fast_decode: int = 0, |
| log_level: int = av.logging.ERROR, overwrite: bool = False) -> bytes: |
| if vcodec not in {"h264", "hevc", "libx264", "libx265", "libsvtav1"}: |
| raise ValueError(f"Unsupported codec {vcodec}") |
| video_path = Path(dst) |
| video_path.parent.mkdir(parents=True, exist_ok=overwrite) |
| if (vcodec in {"libsvtav1", "hevc", "libx265"}) and pix_fmt == "yuv444p": |
| pix_fmt = "yuv420p" |
| h, w, _ = images[0].shape |
| options = {} |
| for k, v in {"g": g, "crf": crf}.items(): |
| if v is not None: |
| options[k] = str(v) |
| if fast_decode: |
| key = "svtav1-params" if vcodec == "libsvtav1" else "tune" |
| options[key] = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode" |
| if log_level is not None: |
| logging.getLogger("libav").setLevel(log_level) |
| with av.open(str(video_path), "w") as out: |
| stream = out.add_stream(vcodec, fps, options=options) |
| stream.pix_fmt, stream.width, stream.height = pix_fmt, w, h |
| for i, img in enumerate(images): |
| frame = av.VideoFrame.from_ndarray(img, format="rgb24") |
| for pkt in stream.encode(frame): |
| out.mux(pkt) |
| if (i + 1) % 100 == 0 or i == len(images) - 1: |
| print(f"Encoded frame {i+1}") |
| for pkt in stream.encode(): |
| out.mux(pkt) |
| if log_level is not None: |
| av.logging.restore_default_callback() |
| if not video_path.exists(): |
| raise OSError(f"Video encoding failed: {video_path}") |
|
|
| def create_video_from_images(images, output_path, fps=30, codec="libx264", quality=23): |
| if not images: |
| raise ValueError("No image data") |
| print(f"Start encoding video, codec: {codec} CRF: {quality}") |
| encode_video_frames(np.asarray(images), Path(output_path), fps=fps, vcodec=codec, crf=quality, overwrite=True) |
| print(f"Video saved to: {output_path}") |
|
|
| |
| def save_data(args, timesteps, actions, dataset_path): |
| data_size = len(actions) |
| data_dict = {k: [] for k in [ |
| '/observations/qpos', '/observations/qvel', '/observations/effort', |
| '/action', '/base_action']} |
|
|
| video_images = {cam: [] for cam in args.camera_names} |
| assert args.export_video, "Please set --export_video to enable video export" |
| while actions: |
| action = actions.pop(0) |
| ts = timesteps.pop(0) |
| for k in ['qpos', 'qvel', 'effort']: |
| data_dict[f'/observations/{k}'].append(ts.observation[k]) |
| data_dict['/action'].append(action) |
| data_dict['/base_action'].append(ts.observation['base_vel']) |
| for cam in args.camera_names: |
| img = ts.observation['images'][cam] |
| if args.export_video: |
| video_img = img if img.shape[2] == 3 else cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
| if video_img.dtype != np.uint8: |
| video_img = (video_img * 255).astype(np.uint8) if video_img.max() <= 1.0 else video_img.astype(np.uint8) |
| video_images[cam].append(video_img) |
|
|
| t0 = time.time() |
| with h5py.File(dataset_path + '.hdf5', 'w', rdcc_nbytes=1024**2*2) as root: |
| root.attrs['sim'], root.attrs['compress'] = False, False |
| obs = root.create_group('observations') |
| for k in ['qpos', 'qvel', 'effort']: |
| obs.create_dataset(k, (data_size, 14)) |
| root.create_dataset('action', (data_size, 14)) |
| root.create_dataset('base_action', (data_size, 2)) |
| for name, arr in data_dict.items(): |
| root[name][...] = arr |
| print(f'\033[32m\nSaving: {time.time() - t0:.1f} secs. %s \033[0m\n' % dataset_path) |
|
|
| if args.export_video and data_size: |
| print('\033[33m\nExporting videos...\033[0m') |
| video_dir = os.path.join(os.path.dirname(dataset_path), "video") |
| os.makedirs(video_dir, exist_ok=True) |
| for cam in args.camera_names: |
| if video_images[cam]: |
| try: |
| cam_dir = os.path.join(video_dir, cam) |
| os.makedirs(cam_dir, exist_ok=True) |
| episode_idx = os.path.basename(dataset_path).split('_')[-1] |
| video_path = os.path.join(cam_dir, f"episode_{episode_idx}.mp4") |
| print(f"Exporting video for camera {cam}: {video_path}") |
| create_video_from_images(video_images[cam], video_path, |
| fps=args.video_fps, codec=args.video_codec, quality=args.video_quality) |
| print(f'\033[32m✅ Camera {cam} video exported: {video_path}\033[0m') |
| except Exception as e: |
| print(f'\033[31m❌ Camera {cam} video export failed: {e}\033[0m') |
| else: |
| print(f'\033[33m⚠️ Camera {cam} has no image data, skipping video export\033[0m') |
| print(f'\033[32m\nVideo export done: {time.time() - t0:.1f} secs\033[0m') |
|
|
| |
| from pynput import keyboard |
|
|
| class RosOperator: |
| def __init__(self, args): |
| self.args = args |
| self.stop_flag = False |
| self.bridge = CvBridge() |
| self.init_deques() |
| self.init_ros() |
|
|
| def init_deques(self): |
| self.img_left_deque = deque() |
| self.img_right_deque = deque() |
| self.img_front_deque = deque() |
| self.img_left_depth_deque = deque() |
| self.img_right_depth_deque = deque() |
| self.img_front_depth_deque = deque() |
| self.master_arm_left_deque = deque() |
| self.master_arm_right_deque = deque() |
| self.puppet_arm_left_deque = deque() |
| self.puppet_arm_right_deque = deque() |
| self.robot_base_deque = deque() |
|
|
| def keyboard_listener(self): |
| def on_press(key): |
| if key == keyboard.Key.space: |
| self.stop_flag = True |
| print("\033[35m>>> Space key pressed, stopping collection and saving...\033[0m") |
| return False |
| with keyboard.Listener(on_press=on_press) as listener: |
| listener.join() |
|
|
| def img_left_callback(self, msg): |
| if len(self.img_left_deque) >= 2000: self.img_left_deque.popleft() |
| self.img_left_deque.append(msg) |
| def img_right_callback(self, msg): |
| if len(self.img_right_deque) >= 2000: self.img_right_deque.popleft() |
| self.img_right_deque.append(msg) |
| def img_front_callback(self, msg): |
| if len(self.img_front_deque) >= 2000: self.img_front_deque.popleft() |
| self.img_front_deque.append(msg) |
| def img_left_depth_callback(self, msg): |
| if len(self.img_left_depth_deque) >= 2000: self.img_left_depth_deque.popleft() |
| self.img_left_depth_deque.append(msg) |
| def img_right_depth_callback(self, msg): |
| if len(self.img_right_depth_deque) >= 2000: self.img_right_depth_deque.popleft() |
| self.img_right_depth_deque.append(msg) |
| def img_front_depth_callback(self, msg): |
| if len(self.img_front_depth_deque) >= 2000: self.img_front_depth_deque.popleft() |
| self.img_front_depth_deque.append(msg) |
| def master_arm_left_callback(self, msg): |
| if len(self.master_arm_left_deque) >= 2000: self.master_arm_left_deque.popleft() |
| self.master_arm_left_deque.append(msg) |
| def master_arm_right_callback(self, msg): |
| if len(self.master_arm_right_deque) >= 2000: self.master_arm_right_deque.popleft() |
| self.master_arm_right_deque.append(msg) |
| def puppet_arm_left_callback(self, msg): |
| if len(self.puppet_arm_left_deque) >= 2000: self.puppet_arm_left_deque.popleft() |
| self.puppet_arm_left_deque.append(msg) |
| def puppet_arm_right_callback(self, msg): |
| if len(self.puppet_arm_right_deque) >= 2000: self.puppet_arm_right_deque.popleft() |
| self.puppet_arm_right_deque.append(msg) |
| def robot_base_callback(self, msg): |
| if len(self.robot_base_deque) >= 2000: self.robot_base_deque.popleft() |
| self.robot_base_deque.append(msg) |
|
|
| def init_ros(self): |
| rospy.init_node('record_episodes', anonymous=True) |
| rospy.Subscriber(self.args.img_left_topic, Image, self.img_left_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_right_topic, Image, self.img_right_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_front_topic, Image, self.img_front_callback, queue_size=1000, tcp_nodelay=True) |
| if self.args.use_depth_image: |
| rospy.Subscriber(self.args.img_left_depth_topic, Image, self.img_left_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_right_depth_topic, Image, self.img_right_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.img_front_depth_topic, Image, self.img_front_depth_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.master_arm_left_topic, JointState, self.master_arm_left_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.master_arm_right_topic, JointState, self.master_arm_right_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.puppet_arm_left_topic, JointState, self.puppet_arm_left_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.puppet_arm_right_topic, JointState, self.puppet_arm_right_callback, queue_size=1000, tcp_nodelay=True) |
| rospy.Subscriber(self.args.robot_base_topic, Odometry, self.robot_base_callback, queue_size=1000, tcp_nodelay=True) |
|
|
| def process(self): |
| timesteps, actions = [], [] |
| count = 0 |
| rate = rospy.Rate(self.args.frame_rate) |
| print("\033[36m>>> Collection started. Press [Space] to stop and save early...\033[0m") |
|
|
| |
| last_qpos_left, last_qpos_right = None, None |
| consecutive_unchanged_count_left, consecutive_unchanged_count_right = 0, 0 |
| UNCHANGED_THRESHOLD = 100 |
|
|
| |
| threading.Thread(target=self.keyboard_listener, daemon=True).start() |
|
|
| while (count < self.args.max_timesteps + 1) and not rospy.is_shutdown() and not self.stop_flag: |
| result = self.get_frame() |
| if not result: |
| rate.sleep() |
| continue |
| count += 1 |
| (img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth, |
| puppet_arm_left, puppet_arm_right, master_arm_left, master_arm_right, robot_base) = result |
|
|
| image_dict = {self.args.camera_names[0]: img_front, |
| self.args.camera_names[1]: img_left, |
| self.args.camera_names[2]: img_right} |
| obs = collections.OrderedDict() |
| obs['images'] = image_dict |
| if self.args.use_depth_image: |
| obs['images_depth'] = {self.args.camera_names[0]: img_front_depth, |
| self.args.camera_names[1]: img_left_depth, |
| self.args.camera_names[2]: img_right_depth} |
| obs['qpos'] = np.concatenate((puppet_arm_left.position, puppet_arm_right.position)) |
| obs['qvel'] = np.concatenate((puppet_arm_left.velocity, puppet_arm_right.velocity)) |
| obs['effort'] = np.concatenate((puppet_arm_left.effort, puppet_arm_right.effort)) |
| obs['base_vel'] = [robot_base.twist.twist.linear.x, robot_base.twist.twist.angular.z] if self.args.use_robot_base else [0.0, 0.0] |
|
|
| |
| current_qpos_left = puppet_arm_left.position |
| if last_qpos_left is not None and np.array_equal(current_qpos_left, last_qpos_left): |
| consecutive_unchanged_count_left += 1 |
| else: |
| consecutive_unchanged_count_left = 0 |
| if consecutive_unchanged_count_left >= UNCHANGED_THRESHOLD: |
| print(f"\033[33m⚠️ Warning: left arm 'position' unchanged for {consecutive_unchanged_count_left} consecutive frames.\033[0m") |
| last_qpos_left = current_qpos_left |
|
|
| |
| current_qpos_right = puppet_arm_right.position |
| if last_qpos_right is not None and np.array_equal(current_qpos_right, last_qpos_right): |
| consecutive_unchanged_count_right += 1 |
| else: |
| consecutive_unchanged_count_right = 0 |
| if consecutive_unchanged_count_right >= UNCHANGED_THRESHOLD: |
| print(f"\033[33m⚠️ Warning: right arm 'position' unchanged for {consecutive_unchanged_count_right} consecutive frames.\033[0m") |
| last_qpos_right = current_qpos_right |
|
|
| if count == 1: |
| timesteps.append(dm_env.TimeStep(dm_env.StepType.FIRST, None, None, obs)) |
| continue |
| timesteps.append(dm_env.TimeStep(dm_env.StepType.MID, None, None, obs)) |
| left_action = np.concatenate((puppet_arm_left.position[:6], [master_arm_left.position[6]])) |
| right_action = np.concatenate((puppet_arm_right.position[:6], [master_arm_right.position[6]])) |
| actions.append(np.concatenate((left_action, right_action))) |
| print("Frame data: ", count) |
| rate.sleep() |
|
|
| print(f"\n>>> Collection finished, {len(actions)} frames total. Saving...") |
| return timesteps, actions |
|
|
| def get_frame(self): |
| if len(self.img_left_deque) == 0 or len(self.img_right_deque) == 0 or len(self.img_front_deque) == 0: |
| return False |
| if self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or len(self.img_right_depth_deque) == 0 or len(self.img_front_depth_deque) == 0): |
| return False |
| frame_time = min([self.img_left_deque[-1].header.stamp.to_sec(), |
| self.img_right_deque[-1].header.stamp.to_sec(), |
| self.img_front_deque[-1].header.stamp.to_sec()]) |
| if self.args.use_depth_image: |
| frame_time = min(frame_time, |
| self.img_left_depth_deque[-1].header.stamp.to_sec(), |
| self.img_right_depth_deque[-1].header.stamp.to_sec(), |
| self.img_front_depth_deque[-1].header.stamp.to_sec()) |
| for dq in [self.img_left_deque, self.img_right_deque, self.img_front_deque, |
| self.master_arm_left_deque, self.master_arm_right_deque, |
| self.puppet_arm_left_deque, self.puppet_arm_right_deque]: |
| if not dq or dq[-1].header.stamp.to_sec() < frame_time: |
| return False |
| if self.args.use_depth_image: |
| for dq in [self.img_left_depth_deque, self.img_right_depth_deque, self.img_front_depth_deque]: |
| if not dq or dq[-1].header.stamp.to_sec() < frame_time: |
| return False |
| if self.args.use_robot_base and (not self.robot_base_deque or self.robot_base_deque[-1].header.stamp.to_sec() < frame_time): |
| return False |
|
|
| def pop(dq): |
| while dq[0].header.stamp.to_sec() < frame_time: |
| dq.popleft() |
| return dq.popleft() |
|
|
| img_left = self.bridge.imgmsg_to_cv2(pop(self.img_left_deque), 'passthrough') |
| img_right = self.bridge.imgmsg_to_cv2(pop(self.img_right_deque), 'passthrough') |
| img_front = self.bridge.imgmsg_to_cv2(pop(self.img_front_deque), 'passthrough') |
| master_arm_left = pop(self.master_arm_left_deque) |
| master_arm_right = pop(self.master_arm_right_deque) |
| puppet_arm_left = pop(self.puppet_arm_left_deque) |
| puppet_arm_right = pop(self.puppet_arm_right_deque) |
| img_left_depth = img_right_depth = img_front_depth = None |
| if self.args.use_depth_image: |
| img_left_depth = cv2.copyMakeBorder(self.bridge.imgmsg_to_cv2(pop(self.img_left_depth_deque), 'passthrough'), 40, 40, 0, 0, cv2.BORDER_CONSTANT, value=0) |
| img_right_depth = cv2.copyMakeBorder(self.bridge.imgmsg_to_cv2(pop(self.img_right_depth_deque), 'passthrough'), 40, 40, 0, 0, cv2.BORDER_CONSTANT, value=0) |
| img_front_depth = cv2.copyMakeBorder(self.bridge.imgmsg_to_cv2(pop(self.img_front_depth_deque), 'passthrough'), 40, 40, 0, 0, cv2.BORDER_CONSTANT, value=0) |
| robot_base = None |
| if self.args.use_robot_base: |
| robot_base = pop(self.robot_base_deque) |
| return (img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth, |
| puppet_arm_left, puppet_arm_right, master_arm_left, master_arm_right, robot_base) |
|
|
| |
| def get_arguments(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--dataset_dir', type=str, default="./data") |
| parser.add_argument('--task_name', type=str, default="aloha_mobile_dummy") |
| parser.add_argument('-e', '--episode_idx', type=int, default=None) |
| parser.add_argument('--max_timesteps', type=int, default=500) |
| parser.add_argument('--camera_names', nargs='+', default=['cam_high', 'cam_left_wrist', 'cam_right_wrist']) |
| parser.add_argument('--img_front_topic', default='/camera_f/color/image_raw') |
| parser.add_argument('--img_left_topic', default='/camera_l/color/image_raw') |
| parser.add_argument('--img_right_topic', default='/camera_r/color/image_raw') |
| parser.add_argument('--img_front_depth_topic', default='/camera_f/depth/image_raw') |
| parser.add_argument('--img_left_depth_topic', default='/camera_l/depth/image_raw') |
| parser.add_argument('--img_right_depth_topic', default='/camera_r/depth/image_raw') |
| parser.add_argument('--master_arm_left_topic', default='/master/joint_left') |
| parser.add_argument('--master_arm_right_topic', default='/master/joint_right') |
| parser.add_argument('--puppet_arm_left_topic', default='/puppet/joint_left') |
| parser.add_argument('--puppet_arm_right_topic', default='/puppet/joint_right') |
| parser.add_argument('--robot_base_topic', default='/odom') |
| parser.add_argument('--use_robot_base', type=bool, default=False) |
| parser.add_argument('--use_depth_image', type=bool, default=False) |
| parser.add_argument('--frame_rate', type=int, default=30) |
| parser.add_argument('--export_video', action='store_true', help='Enable video export') |
| parser.add_argument('--video_fps', type=int, default=30) |
| parser.add_argument('--video_codec', choices=['libx264', 'libx265', 'libsvtav1'], default='libx264') |
| parser.add_argument('--video_quality', type=int, default=23, help='CRF value; lower = higher quality') |
|
|
| args = parser.parse_args() |
|
|
| |
| if args.episode_idx is None: |
| dataset_dir = os.path.join(args.dataset_dir, args.task_name) |
| episode_indices = [] |
| if os.path.exists(dataset_dir) and os.path.isdir(dataset_dir): |
| pattern = re.compile(r'^episode_(\d+)\.hdf5$') |
| for filename in os.listdir(dataset_dir): |
| match = pattern.match(filename) |
| if match: |
| episode_indices.append(int(match.group(1))) |
| if episode_indices: |
| args.episode_idx = max(episode_indices) + 1 |
| print(f"\033[34m>>> Found {len(episode_indices)} existing episodes, next index: {args.episode_idx}\033[0m") |
| else: |
| args.episode_idx = 0 |
| print(f"\033[34m>>> No existing episodes found, starting from index 0\033[0m") |
|
|
| return args |
|
|
| |
| def main(): |
| args = get_arguments() |
| ros_operator = RosOperator(args) |
| timesteps, actions = ros_operator.process() |
|
|
| if len(actions) == 0: |
| print("\033[31m\nNo data collected, skipping save.\033[0m") |
| return |
|
|
| dataset_dir = os.path.join(args.dataset_dir, args.task_name) |
| os.makedirs(dataset_dir, exist_ok=True) |
| dataset_path = os.path.join(dataset_dir, f"episode_{args.episode_idx}") |
| save_data(args, timesteps, actions, dataset_path) |
| print("\033[32m>>> Saved to:", dataset_path + ".hdf5\033[0m") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
