collect_data/docs/collect_data.py

#coding=utf-8
import os
import time
import numpy as np
import h5py
import argparse
import dm_env
import cv2
import queue

import rospy
from sensor_msgs.msg import JointState
from sensor_msgs.msg import Image
from nav_msgs.msg import Odometry
from cv_bridge import CvBridge

import threading
import collections

import sys
sys.path.append("./")



# 保存数据函数
def save_data(args, timesteps, actions, dataset_path):
    use_depth_image = args.use_depth_image
    is_compress = args.is_compress
    
    # 数据字典
    data_dict = {
        # obs包含qpos，qvel， effort ,acition, image, base_action
        '/observations/qpos': [],
        '/observations/qvel': [],
        '/observations/effort': [],
        '/action': [],
        '/base_action': [],
    }

    # 相机字典  观察的图像
    for cam_name in args.camera_names:
        data_dict[f'/observations/images/{cam_name}'] = []
        if use_depth_image:
             data_dict[f'/observations/depths/{cam_name}'] = []
        

    # len(action): max_timesteps, len(time_steps): max_timesteps + 1
    # 动作长度 遍历动作
    while actions:
        # 循环弹出一个队列
        action = actions.pop(0)   # 动作  当前动作
        ts = timesteps.pop(0)     # 奖励  前一帧

        # 往字典里面添值
        # Timestep返回的qpos，qvel,effort
        data_dict['/observations/qpos'].append(ts.observation['qpos'])
        data_dict['/observations/qvel'].append(ts.observation['qvel'])
        data_dict['/observations/effort'].append(ts.observation['effort'])

        # 实际发的action
        data_dict['/action'].append(action)
        data_dict['/base_action'].append(ts.observation['base_vel'])

        # 相机数据
        for cam_name in args.camera_names:
            data_dict[f'/observations/images/{cam_name}'].append(ts.observation['images'][cam_name])
            
            if use_depth_image:
                data_dict[f'/observations/depths/{cam_name}'].append(ts.observation['depths'][cam_name])

    pass
    if is_compress:
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 50] # tried as low as 20, seems fine
        compressed_len = []
        # 3个相机
        for cam_name in args.camera_names:
            image_list = data_dict[f'/observations/images/{cam_name}']
            compressed_list = []
            compressed_len.append([])   # 压缩的长度
            
            for image in image_list:
                result, encoded_image = cv2.imencode('.jpg', image, encode_param) # 0.02 sec # cv2.imdecode(encoded_image, 1)
                
                compressed_list.append(encoded_image)
                compressed_len[-1].append(len(encoded_image))
            
            # 更新图像
            data_dict[f'/observations/images/{cam_name}'] = compressed_list

        compressed_len = np.array(compressed_len)
        padded_size = compressed_len.max()  # 取最大的图像长度，图像压缩后就是一个buf序列
        
        for cam_name in args.camera_names:
            compressed_image_list = data_dict[f'/observations/images/{cam_name}']
            padded_compressed_image_list = []
            for compressed_image in compressed_image_list:
                padded_compressed_image = np.zeros(padded_size, dtype='uint8')
                image_len = len(compressed_image)
                padded_compressed_image[:image_len] = compressed_image
                padded_compressed_image_list.append(padded_compressed_image)
            data_dict[f'/observations/images/{cam_name}'] = padded_compressed_image_list

        # 深度图压缩处理
        if use_depth_image:
            compressed_len_depth = []
            # 3个相机
            for cam_name in args.camera_names:
                depth_list = data_dict[f'/observations/depths/{cam_name}']
                compressed_list_depth = []
                compressed_len_depth.append([])   # 压缩的长度
                
                for depth in depth_list:
                    result, encoded_depth = cv2.imencode('.jpg', depth, encode_param) # 0.02 sec # cv2.imdecode(encoded_image, 1)
                    
                    compressed_list_depth.append(encoded_depth)
                    compressed_len_depth[-1].append(len(encoded_depth))
                
                # 更新图像
                data_dict[f'/observations/depths/{cam_name}'] = compressed_list_depth

            compressed_len_depth = np.array(compressed_len_depth)
            padded_size_depth = compressed_len_depth.max()  # 取最大的图像长度，图像压缩后就是一个buf序列
            
            for cam_name in args.camera_names:
                compressed_depth_list = data_dict[f'/observations/depths/{cam_name}']
                padded_compressed_depth_list = []
                for compressed_depth in compressed_depth_list:
                    padded_compressed_depth = np.zeros(padded_size_depth, dtype='uint8')
                    depth_len = len(compressed_depth)
                    padded_compressed_depth[:depth_len] = compressed_depth
                    padded_compressed_depth_list.append(padded_compressed_depth)
                data_dict[f'/observations/depths/{cam_name}'] = padded_compressed_depth_list




    t0 = time.time()
    with h5py.File(dataset_path + '.hdf5', 'w', rdcc_nbytes=1024**2*2) as root:
        # 文本的属性：
        # 1 是否仿真
        # 2 图像是否压缩
        #
        root.attrs['sim'] = False
        root.attrs['compress'] = False
        if args.is_compress:
            root.attrs['compress'] = True

        # 创建一个新的组observations，观测状态组
        # 图像组
        obs = root.create_group('observations')
        
        image = obs.create_group('images')

        depth = obs.create_group('depths')

        for cam_name in args.camera_names:
            if args.is_compress:
                _ = image.create_dataset(cam_name, (args.max_timesteps, padded_size), dtype='uint8',
                                         chunks=(1, padded_size), )        
                if use_depth_image:
                    # _ = depth.create_dataset(cam_name, (args.max_timesteps, 400, 640), dtype='uint8',
                    #                      chunks=(1, 400, 640), )
                    _ = depth.create_dataset(cam_name, (args.max_timesteps, padded_size_depth), dtype='uint8',
                                         chunks=(1, padded_size_depth), )       


            else:
                _ = image.create_dataset(cam_name, (args.max_timesteps, 480, 640, 3), dtype='uint8',
                                         chunks=(1, 480, 640, 3), )
                if use_depth_image:
                    _ = depth.create_dataset(cam_name, (args.max_timesteps, 400, 640), dtype='uint8',
                                         chunks=(1, 400, 640), )

        _ = obs.create_dataset('qpos', (args.max_timesteps, 14))
        _ = obs.create_dataset('qvel', (args.max_timesteps, 14))
        _ = obs.create_dataset('effort', (args.max_timesteps, 14))
        _ = root.create_dataset('action', (args.max_timesteps, 14))
        _ = root.create_dataset('base_action', (args.max_timesteps, 2))


        if is_compress:
            _ = root.create_dataset('compress_len', (len(args.camera_names), args.max_timesteps))
            root['/compress_len'][...] = compressed_len

            
        # data_dict写入h5py.File
        for name, array in data_dict.items():   # 名字+值
            root[name][...] = array

    print(f'Saving: {time.time() - t0:.1f} secs', dataset_path)


class RosOperator:
    def __init__(self, args):
        self.init(args)

    def init(self, args):
        self.args = args
        self.use_depth_image = self.args.use_depth_image
        self.use_robot_base  = self.args.use_robot_base

        self.bridge = CvBridge()
        self.imgl_queue, self.imgr_queue, self.imgf_queue = [queue.Queue() for _ in range(3)]
        self.depthl_queue, self.depthr_queue, self.depthf_queue = [queue.Queue() for _ in range(3)]
        self.masterl_queue, self.masterr_queue, self.puppetl_queue, self.puppetr_queue = [queue.Queue() for _ in range(4)]
        self.robot_base_deque = queue.Queue()

        dataset_dir = os.path.join(args.dataset_dir, args.task_name)
        if not os.path.exists(dataset_dir):
            os.makedirs(dataset_dir)
        self.dataset_path = os.path.join(dataset_dir, "episode_" + str(args.episode_idx))

        self.register_sub()   # 定义订阅消息话题

    def register_sub(self):
        rospy.init_node('record_episodes', anonymous=True)
        
        rospy.Subscriber(self.args.img_left_topic,  Image, lambda msg: self.imgl_queue.put(msg), queue_size=60, tcp_nodelay=True)
        rospy.Subscriber(self.args.img_right_topic, Image, lambda msg: self.imgr_queue.put(msg), queue_size=60, tcp_nodelay=True)
        rospy.Subscriber(self.args.img_front_topic, Image, lambda msg: self.imgf_queue.put(msg), queue_size=60, tcp_nodelay=True)
        
        if self.use_depth_image:
            rospy.Subscriber(self.args.depth_left_topic,  Image, lambda msg: self.depthl_queue.put(msg), queue_size=60, tcp_nodelay=True)
            rospy.Subscriber(self.args.depth_right_topic, Image, lambda msg: self.depthr_queue.put(msg), queue_size=60, tcp_nodelay=True)
            rospy.Subscriber(self.args.depth_front_topic, Image, lambda msg: self.depthf_queue.put(msg), queue_size=60, tcp_nodelay=True)


        rospy.Subscriber(self.args.master_arm_left_topic,  JointState, lambda msg: self.masterl_queue.put(msg), queue_size=200, tcp_nodelay=True)
        rospy.Subscriber(self.args.master_arm_right_topic, JointState, lambda msg: self.masterr_queue.put(msg), queue_size=200, tcp_nodelay=True)
        rospy.Subscriber(self.args.puppet_arm_left_topic,  JointState, lambda msg: self.puppetl_queue.put(msg), queue_size=200, tcp_nodelay=True)
        rospy.Subscriber(self.args.puppet_arm_right_topic, JointState, lambda msg: self.puppetr_queue.put(msg), queue_size=200, tcp_nodelay=True)
        
        if self.use_robot_base:
            rospy.Subscriber(self.args.robot_base_topic, Odometry, lambda msg: self.robot_base_deque.put(msg), queue_size=60, tcp_nodelay=True)

 
    def collet_data(self):
        rate = rospy.Rate(30)
        input("\033[32m\nPlease press any key to continue\n\033[0m")
        
        count = 0      # 统计 timesteps
        timesteps = [] # obs
        actions = []   # acticn
        
        # 循环max_timesteps次
        while (count < self.args.max_timesteps + 1 and not rospy.is_shutdown()):
            time1 = time.time()     # 起始时间
            # 判读图像和深度图队列是否为空
            if self.imgl_queue.empty() or self.imgr_queue.empty() or self.imgf_queue.empty() or \
                (self.use_depth_image and (self.depthl_queue.empty() or self.depthr_queue.empty() or self.depthf_queue.empty())):
                print("\033[31mPlease check topic\n\033[0m")
                rate.sleep()
                continue
            
            if self.masterl_queue.empty() or self.masterr_queue.empty() or \
                self.puppetl_queue.empty() or self.puppetr_queue.empty() or \
                (  self.use_robot_base and ( self.robot_base_deque.empty() )  ):
                print("\033[31mPlease check topic\n\033[0m")
                rate.sleep()
                continue
            
            imgl, imgr, imgl, depthl, depthr, depthl = [None for _ in range(6)]
            masterl, masterr, puppetl, puppetr = [None for _ in range(4)]
            robot_base = None

            min_timestamps = min(self.imgl_queue.queue[-1].header.stamp.to_sec(),
                                 self.imgr_queue.queue[-1].header.stamp.to_sec(),
                                 self.imgf_queue.queue[-1].header.stamp.to_sec())  - 0.01
            
            if(self.use_depth_image):
                min_timestamps = min(min_timestamps, 
                                     self.depthl_queue.queue[-1].header.stamp.to_sec(),
                                     self.depthr_queue.queue[-1].header.stamp.to_sec(),
                                     self.depthf_queue.queue[-1].header.stamp.to_sec(),
                                    ) - 0.01

            while (self.imgl_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.imgl_queue.get()
            imgl = self.bridge.imgmsg_to_cv2(self.imgl_queue.get(), 'passthrough')

            while (self.imgr_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.imgr_queue.get()
            imgr = self.bridge.imgmsg_to_cv2(self.imgr_queue.get(), 'passthrough')

            while (self.imgf_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.imgf_queue.get()
            imgf = self.bridge.imgmsg_to_cv2(self.imgf_queue.get(), 'passthrough')
            
            if(self.use_depth_image):
                while (self.depthl_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                    self.depthl_queue.get()
                depthl = self.bridge.imgmsg_to_cv2(self.depthl_queue.get(), 'passthrough')

                while (self.depthr_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                    self.depthr_queue.get()
                depthr = self.bridge.imgmsg_to_cv2(self.depthr_queue.get(), 'passthrough')

                while (self.depthf_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                    self.depthf_queue.get()
                depthf = self.bridge.imgmsg_to_cv2(self.depthf_queue.get(), 'passthrough')
            


            while (self.masterl_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.masterl_queue.get()
            masterl = self.masterl_queue.get()

            while (self.masterr_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.masterr_queue.get()
            masterr = self.masterr_queue.get()

            while (self.puppetl_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.puppetl_queue.get()
            puppetl = self.puppetl_queue.get()

            while (self.puppetr_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.puppetr_queue.get()
            puppetr = self.puppetl_queue.get()

            if(self.use_robot_base):
                while self.robot_base_deque[0].header.stamp.to_sec() < min_timestamps:
                    self.robot_base_deque.get()
                robot_base = self.robot_base_deque.get()

            # cv2.imshow("imgl", imgl)
            # cv2.imshow("imgr", imgr)
            # cv2.imshow("imgf", imgf)
            # cv2.waitKey(20)

            count += 1
           

            # 2.1 图像信息
            image_dict = dict()
            image_dict[self.args.camera_names[0]] = imgf
            image_dict[self.args.camera_names[1]] = imgl
            image_dict[self.args.camera_names[2]] = imgr

            # 2.2 从臂的信息从臂的状态 机械臂示教模式时 会自动订阅
            obs = collections.OrderedDict()  # 有序的字典
            obs['images'] = image_dict
            obs['qpos'] = np.concatenate((np.array(puppetl.position), np.array(puppetr.position)), axis=0)
            obs['qvel'] = np.concatenate((np.array(puppetl.velocity), np.array(puppetr.velocity)), axis=0)
            obs['effort'] = np.concatenate((np.array(puppetl.effort), np.array(puppetr.effort)), axis=0)
            
            # 2.3 底盘数据
            if self.use_robot_base:
                obs['base_vel'] = [robot_base.twist.twist.linear.x, robot_base.twist.twist.angular.z]
            else:
                obs['base_vel'] = [0.0, 0.0]

            # 2.4 深度图数据
            if(self.use_depth_image):
                depth_dict = dict()
                depth_dict[self.args.camera_names[0]] = depthf
                depth_dict[self.args.camera_names[1]] = depthl
                depth_dict[self.args.camera_names[2]] = depthr
                obs['depths'] = depth_dict

            # 第一帧 只包含first， fisrt只保存StepType.FIRST
            if count == 1:
                ts = dm_env.TimeStep(
                    step_type=dm_env.StepType.FIRST,
                    reward=None,
                    discount=None,
                    observation=obs)
                
                timesteps.append(ts)
                print("frame %s cost_time: %s" % (count, round(1/(time.time() - time1), 6)))
                rate.sleep()
                continue

            # 时间步
            ts = dm_env.TimeStep(
                step_type=dm_env.StepType.MID,
                reward=None,
                discount=None,
                observation=obs)

            # 主臂保存状态
            action = np.concatenate((np.array(masterl.position), np.array(masterr.position)), axis=0)
            actions.append(action)
            timesteps.append(ts)

            if rospy.is_shutdown():
                exit(-1)
            
            print("frame %s cost_time: %s" % (count, round((time.time() - time1), 6)))
            rate.sleep()
            
        print("len(timesteps): ", len(timesteps))
        print("len(actions)  : ", len(actions))
        save_data(self.args, timesteps, actions, self.dataset_path)

def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_dir', action='store', type=str, help='Dataset_dir.',
                        default="./data", required=False)
    
    parser.add_argument('--task_name', action='store', type=str, help='Task name.',
                        default="aloha_mobile_dummy", required=False)
    
    parser.add_argument('--episode_idx', action='store', type=int, help='Episode index.',
                        default=0, required=False)
    parser.add_argument('--max_timesteps', action='store', type=int, help='Max_timesteps.',
                        default=500, required=False)

    # 相机名称话题
    parser.add_argument('--camera_names', action='store', type=str, help='camera_names',
                        default=['cam_high', 'cam_left_wrist', 'cam_right_wrist'], required=False)
    
    # 相机彩色图话题
    parser.add_argument('--img_front_topic', action='store', type=str, help='img_front_topic',
                        default='/camera_f/color/image_raw', required=False)
    parser.add_argument('--img_left_topic', action='store', type=str, help='img_left_topic',
                        default='/camera_l/color/image_raw', required=False)
    parser.add_argument('--img_right_topic', action='store', type=str, help='img_right_topic',
                        default='/camera_r/color/image_raw', required=False)
    
    # 相机深度图话题
    parser.add_argument('--use_depth_image', action='store_true', help='--use_depth_image', required=False)
    
    parser.add_argument('--depth_front_topic', action='store', type=str, help='depth_front_topic',
                        default='/camera_f/depth/image_raw', required=False)
    parser.add_argument('--depth_left_topic', action='store', type=str, help='depth_left_topic',
                        default='/camera_l/depth/image_raw', required=False)
    parser.add_argument('--depth_right_topic', action='store', type=str, help='depth_right_topic',
                        default='/camera_r/depth/image_raw', required=False)

    # 机械臂话题
    parser.add_argument('--master_arm_left_topic', action='store', type=str, help='master_arm_left_topic',
                        default='/master/joint_left', required=False)
    parser.add_argument('--master_arm_right_topic', action='store', type=str, help='master_arm_right_topic',
                        default='/master/joint_right', required=False)
    
    parser.add_argument('--puppet_arm_left_topic', action='store', type=str, help='puppet_arm_left_topic',
                        default='/puppet/joint_left', required=False)
    parser.add_argument('--puppet_arm_right_topic', action='store', type=str, help='puppet_arm_right_topic',
                        default='/puppet/joint_right', required=False)
    # 底盘话题
    parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base',
                        default=False, required=False)
    parser.add_argument('--robot_base_topic', action='store', type=str, help='robot_base_topic',
                        default='/odom_raw', required=False)
    
    # 发布频率
    parser.add_argument('--frame_rate', action='store', type=int, help='frame_rate',
                        default=30, required=False)
    
    # 图像是否压缩
    parser.add_argument('--is_compress', action='store_true', help='is_compress', required=False)
    
    args = parser.parse_args()
    return args


def main():
    args = get_arguments()
    ros_operator = RosOperator(args)   # 初始化ros节点，订阅消息

    t1 = threading.Thread(target=ros_operator.collet_data)
    t1.start()
    t1.join()
    

if __name__ == '__main__':
    main()

# python scripts/record_data.py  --dataset_dir ~/data0301 --max_timesteps 500 --is_compress --episode_idx 0