#!/usr/bin/env python # -- coding: UTF-8 -- 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}") # ── save_data ───────────────────────────────────────────────────────────────── 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') # ── RosOperator ─────────────────────────────────────────────────────────────── 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") # qpos change detection variables (left/right arm separately) last_qpos_left, last_qpos_right = None, None consecutive_unchanged_count_left, consecutive_unchanged_count_right = 0, 0 UNCHANGED_THRESHOLD = 100 # warn after 100 consecutive unchanged frames # start background keyboard listener thread 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] # left arm qpos change detection 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 # right arm qpos change detection 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) # ── Argument parsing ────────────────────────────────────────────────────────── 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() # auto-compute episode_idx if not specified 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 # ── main ────────────────────────────────────────────────────────────────────── 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()