306 lines
15 KiB
Python
306 lines
15 KiB
Python
import cv2
|
||
import numpy as np
|
||
import dm_env
|
||
|
||
import collections
|
||
from collections import deque
|
||
|
||
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
|
||
from utils import display_camera_grid, save_data
|
||
|
||
|
||
class AlohaRobotRos:
|
||
def __init__(self, args):
|
||
self.robot_base_deque = None
|
||
self.puppet_arm_right_deque = None
|
||
self.puppet_arm_left_deque = None
|
||
self.master_arm_right_deque = None
|
||
self.master_arm_left_deque = None
|
||
self.img_front_deque = None
|
||
self.img_right_deque = None
|
||
self.img_left_deque = None
|
||
self.img_front_depth_deque = None
|
||
self.img_right_depth_deque = None
|
||
self.img_left_depth_deque = None
|
||
self.bridge = None
|
||
self.args = args
|
||
self.init()
|
||
self.init_ros()
|
||
|
||
def init(self):
|
||
self.bridge = CvBridge()
|
||
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 get_frame(self):
|
||
print(len(self.img_left_deque), len(self.img_right_deque), len(self.img_front_deque),
|
||
len(self.img_left_depth_deque), len(self.img_right_depth_deque), len(self.img_front_depth_deque))
|
||
if len(self.img_left_deque) == 0 or len(self.img_right_deque) == 0 or len(self.img_front_deque) == 0 or \
|
||
(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
|
||
if self.args.use_depth_image:
|
||
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(),
|
||
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()])
|
||
else:
|
||
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 len(self.img_left_deque) == 0 or self.img_left_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.img_right_deque) == 0 or self.img_right_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.img_front_deque) == 0 or self.img_front_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.master_arm_left_deque) == 0 or self.master_arm_left_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.master_arm_right_deque) == 0 or self.master_arm_right_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.puppet_arm_left_deque) == 0 or self.puppet_arm_left_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if len(self.puppet_arm_right_deque) == 0 or self.puppet_arm_right_deque[-1].header.stamp.to_sec() < frame_time:
|
||
return False
|
||
if self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or self.img_left_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
||
return False
|
||
if self.args.use_depth_image and (len(self.img_right_depth_deque) == 0 or self.img_right_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
||
return False
|
||
if self.args.use_depth_image and (len(self.img_front_depth_deque) == 0 or self.img_front_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
||
return False
|
||
if self.args.use_robot_base and (len(self.robot_base_deque) == 0 or self.robot_base_deque[-1].header.stamp.to_sec() < frame_time):
|
||
return False
|
||
|
||
while self.img_left_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_left_deque.popleft()
|
||
img_left = self.bridge.imgmsg_to_cv2(self.img_left_deque.popleft(), 'passthrough')
|
||
# print("img_left:", img_left.shape)
|
||
|
||
while self.img_right_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_right_deque.popleft()
|
||
img_right = self.bridge.imgmsg_to_cv2(self.img_right_deque.popleft(), 'passthrough')
|
||
|
||
while self.img_front_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_front_deque.popleft()
|
||
img_front = self.bridge.imgmsg_to_cv2(self.img_front_deque.popleft(), 'passthrough')
|
||
|
||
while self.master_arm_left_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.master_arm_left_deque.popleft()
|
||
master_arm_left = self.master_arm_left_deque.popleft()
|
||
|
||
while self.master_arm_right_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.master_arm_right_deque.popleft()
|
||
master_arm_right = self.master_arm_right_deque.popleft()
|
||
|
||
while self.puppet_arm_left_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.puppet_arm_left_deque.popleft()
|
||
puppet_arm_left = self.puppet_arm_left_deque.popleft()
|
||
|
||
while self.puppet_arm_right_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.puppet_arm_right_deque.popleft()
|
||
puppet_arm_right = self.puppet_arm_right_deque.popleft()
|
||
|
||
img_left_depth = None
|
||
if self.args.use_depth_image:
|
||
while self.img_left_depth_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_left_depth_deque.popleft()
|
||
img_left_depth = self.bridge.imgmsg_to_cv2(self.img_left_depth_deque.popleft(), 'passthrough')
|
||
top, bottom, left, right = 40, 40, 0, 0
|
||
img_left_depth = cv2.copyMakeBorder(img_left_depth, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
|
||
|
||
img_right_depth = None
|
||
if self.args.use_depth_image:
|
||
while self.img_right_depth_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_right_depth_deque.popleft()
|
||
img_right_depth = self.bridge.imgmsg_to_cv2(self.img_right_depth_deque.popleft(), 'passthrough')
|
||
top, bottom, left, right = 40, 40, 0, 0
|
||
img_right_depth = cv2.copyMakeBorder(img_right_depth, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
|
||
|
||
img_front_depth = None
|
||
if self.args.use_depth_image:
|
||
while self.img_front_depth_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.img_front_depth_deque.popleft()
|
||
img_front_depth = self.bridge.imgmsg_to_cv2(self.img_front_depth_deque.popleft(), 'passthrough')
|
||
top, bottom, left, right = 40, 40, 0, 0
|
||
img_front_depth = cv2.copyMakeBorder(img_front_depth, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
|
||
|
||
robot_base = None
|
||
if self.args.use_robot_base:
|
||
while self.robot_base_deque[0].header.stamp.to_sec() < frame_time:
|
||
self.robot_base_deque.popleft()
|
||
robot_base = self.robot_base_deque.popleft()
|
||
|
||
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 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):
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
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 = []
|
||
# 图像数据
|
||
image = np.random.randint(0, 255, size=(480, 640, 3), dtype=np.uint8)
|
||
image_dict = dict()
|
||
for cam_name in self.args.camera_names:
|
||
image_dict[cam_name] = image
|
||
count = 0
|
||
|
||
# input_key = input("please input s:")
|
||
# while input_key != 's' and not rospy.is_shutdown():
|
||
# input_key = input("please input s:")
|
||
|
||
rate = rospy.Rate(self.args.frame_rate)
|
||
print_flag = True
|
||
|
||
while (count < self.args.max_timesteps + 1) and not rospy.is_shutdown():
|
||
# 2 收集数据
|
||
result = self.get_frame()
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
if not result:
|
||
if print_flag:
|
||
print("syn fail\n")
|
||
print_flag = False
|
||
rate.sleep()
|
||
continue
|
||
print_flag = True
|
||
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
|
||
# 2.1 图像信息
|
||
image_dict = dict()
|
||
image_dict[self.args.camera_names[0]] = img_front
|
||
image_dict[self.args.camera_names[1]] = img_left
|
||
image_dict[self.args.camera_names[2]] = img_right
|
||
|
||
# import pdb
|
||
# pdb.set_trace()
|
||
display_camera_grid(image_dict)
|
||
|
||
# 2.2 从臂的信息从臂的状态 机械臂示教模式时 会自动订阅
|
||
obs = collections.OrderedDict() # 有序的字典
|
||
obs['images'] = image_dict
|
||
if self.args.use_depth_image:
|
||
image_dict_depth = dict()
|
||
image_dict_depth[self.args.camera_names[0]] = img_front_depth
|
||
image_dict_depth[self.args.camera_names[1]] = img_left_depth
|
||
image_dict_depth[self.args.camera_names[2]] = img_right_depth
|
||
obs['images_depth'] = image_dict_depth
|
||
obs['qpos'] = np.concatenate((np.array(puppet_arm_left.position), np.array(puppet_arm_right.position)), axis=0)
|
||
obs['qvel'] = np.concatenate((np.array(puppet_arm_left.velocity), np.array(puppet_arm_right.velocity)), axis=0)
|
||
obs['effort'] = np.concatenate((np.array(puppet_arm_left.effort), np.array(puppet_arm_right.effort)), axis=0)
|
||
if self.args.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]
|
||
|
||
# 第一帧 只包含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)
|
||
continue
|
||
|
||
# 时间步
|
||
ts = dm_env.TimeStep(
|
||
step_type=dm_env.StepType.MID,
|
||
reward=None,
|
||
discount=None,
|
||
observation=obs)
|
||
|
||
# 主臂保存状态
|
||
action = np.concatenate((np.array(master_arm_left.position), np.array(master_arm_right.position)), axis=0)
|
||
actions.append(action)
|
||
timesteps.append(ts)
|
||
print(f"\n{self.args.episode_idx} | Frame data: {count}\r", end="")
|
||
if rospy.is_shutdown():
|
||
exit(-1)
|
||
rate.sleep()
|
||
|
||
print("len(timesteps): ", len(timesteps))
|
||
print("len(actions) : ", len(actions))
|
||
return timesteps, actions
|