import cv2
import numpy as np
import zmq
import time
import threading
from enum import Enum
import os
# 定义机器人观测数据结构
class RobotObsShape:
HEAD_IMAGE_SHAPE = (480, 640, 3)
HEAD_IMAGE_SIZE = HEAD_IMAGE_SHAPE[0] * HEAD_IMAGE_SHAPE[1] * HEAD_IMAGE_SHAPE[2]
WRIST_IMAGE_SHAPE = (480, 640, 3)
WRIST_IMAGE_SIZE = WRIST_IMAGE_SHAPE[0] * WRIST_IMAGE_SHAPE[1] * WRIST_IMAGE_SHAPE[2]
STATE_SHAPE = (7,)
STATE_SIZE = STATE_SHAPE[0] * 4 # float32占4字节
CHUNK_SIZE = HEAD_IMAGE_SIZE + WRIST_IMAGE_SIZE + STATE_SIZE
DOF = 7 # 自由度
NUM_ACTIONS = 8 # 动作数量
ACTIONS_SHAPE = (NUM_ACTIONS, DOF)
ACTIONS_SIZE = NUM_ACTIONS * DOF * 4 # 总动作数据大小
# 机器人状态枚举
class RobotState(str, Enum):
IDLE = "IDLE"
INITALIZING = "INITALIZING"
ACTING = "ACTING"
FINISHED = "FINISHED"
def create_test_image(height, width):
"""创建测试图像:彩色渐变"""
img = np.zeros((height, width, 3), dtype=np.uint8)
for i in range(height):
for j in range(width):
img[i, j] = [
int(255 * j / width), # 红色通道:从左到右渐变
int(255 * i / height), # 绿色通道:从上到下渐变
150 # 蓝色通道:固定值
]
return img
class Client:
def __init__(self, port=15558):
self.context = zmq.Context()
self.socket = self.context.socket(zmq.REQ)
self.socket.connect(f"tcp://192.168.2.211:{port}")
self.socket.RCVTIMEO = 1000 # 设置接收超时为1秒
self.running = True
self.state = RobotState.IDLE
print(f"Client connected to port {port}")
# # 创建测试图像
# self.head_img = create_test_image(*RobotObsShape.HEAD_IMAGE_SHAPE[:2])
# self.wrist_img = create_test_image(*RobotObsShape.WRIST_IMAGE_SHAPE[:2])
# # 初始关节角度 (6个关节角度 + 夹爪状态)
# self.joint_angles = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
self.head_img = cv2.imread("/home/fishros/hdx/tool/dataset_ori/imgs/0_0.jpg")
self.wrist_img = cv2.imread("/home/fishros/hdx/tool/dataset_ori/imgs/0_0.jpg")
# 初始关节角度
self.joint_angles = np.array([244.021978021978,39.32967032967033,17.208791208791208,291.4725274725275,119.56043956043956,75.05494505494505,0], dtype=np.float32)
# 启动键盘监听线程
# self.listening_thread = threading.Thread(target=self.listen_keyboard)
# self.listening_thread.daemon = True
# self.listening_thread.start()
# 创建保存图像的目录
# os.makedirs("saved_images", exist_ok=True)
# def listen_keyboard(self):
# """监听键盘输入线程"""
# print("Keyboard listener started. Commands: [q]退出, [s]保存图片, [a]发送消息,[i]初始化关节")
# while self.running:
# cmd = input("> ").lower()
# if cmd == 'q':
# self.running = False
# print("Exiting...")
# elif cmd == 's':
# # 保存当前图像
# cv2.imwrite("saved_images/head_image.jpg", self.head_img)
# cv2.imwrite("saved_images/wrist_image.jpg", self.wrist_img)
# print("Images saved to saved_images/ directory")
# elif cmd == 'a':
# self.state = RobotState.ACTING
# print("State changed to ACTING")
# elif cmd == 'i':
# self.state = RobotState.INITALIZING
# print("State changed to INITIALIZING")
# elif cmd == 'f':
# self.state = RobotState.FINISHED
# print("State changed to FINISHED")
def create_message(self):
"""创建符合协议的消息字节流"""
# 确保图像尺寸正确
if self.head_img.shape != RobotObsShape.HEAD_IMAGE_SHAPE:
self.head_img = cv2.resize(self.head_img,
(RobotObsShape.HEAD_IMAGE_SHAPE[1],
RobotObsShape.HEAD_IMAGE_SHAPE[0]))
if self.wrist_img.shape != RobotObsShape.WRIST_IMAGE_SHAPE:
self.wrist_img = cv2.resize(self.wrist_img,
(RobotObsShape.WRIST_IMAGE_SHAPE[1],
RobotObsShape.WRIST_IMAGE_SHAPE[0]))
# 创建消息缓冲区
buffer = bytearray(RobotObsShape.CHUNK_SIZE)
# 填充头部图像数据
head_bytes = self.head_img.tobytes()
buffer[:len(head_bytes)] = head_bytes
# 填充腕部图像数据
start = RobotObsShape.HEAD_IMAGE_SIZE
end = start + RobotObsShape.WRIST_IMAGE_SIZE
wrist_bytes = self.wrist_img.tobytes()
buffer[start:end] = wrist_bytes
# 填充关节角度数据
start = end
end = start + RobotObsShape.STATE_SIZE
buffer[start:end] = self.joint_angles.tobytes()
return bytes(buffer)
def parse_actions(self, action_bytes):
"""解析接收到的动作数据"""
# 验证消息长度
if len(action_bytes) != RobotObsShape.ACTIONS_SHAPE:
return None
actions = np.frombuffer(action_bytes, dtype=np.float32)
return actions.reshape(RobotObsShape.ACTIONS_SHAPE)
def run(self):
"""主循环:发送观察数据并接收动作"""
iteration = 0
while self.running:
# 在图像上添加状态信息
# status_text = f"Iteration: {iteration} | State: {self.state}"
# cv2.putText(self.head_img, status_text, (10, 30),
# cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
# 显示图像
cv2.imshow("Head Camera", self.head_img)
#cv2.imshow("Wrist Camera", self.wrist_img)
# 处理按键事件(关键修改:放在主循环中)
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
self.running = False
break
elif key == ord('a'):
self.state = RobotState.ACTING
print("State changed to ACTING")
elif key == ord('i'):
self.state = RobotState.INITALIZING
print("State changed to INITIALIZING")
elif key == ord('f'):
self.state = RobotState.FINISHED
print("State changed to FINISHED")
elif key == ord('s'):
cv2.imwrite("saved_images/head_image.jpg", self.head_img)
cv2.imwrite("saved_images/wrist_image.jpg", self.wrist_img)
print("Images saved to saved_images/ directory")
# 状态处理
if self.state == RobotState.INITALIZING:
# 初始关节角度 (6个关节角度 + 夹爪状态)
#self.joint_angles = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
# 创建新的测试图像
#self.head_img = create_test_image(*RobotObsShape.HEAD_IMAGE_SHAPE[:2])
#self.wrist_img = create_test_image(*RobotObsShape.WRIST_IMAGE_SHAPE[:2])
self.head_img = cv2.imread("/home/fishros/hdx/tool/dataset_ori/imgs/0_0.jpg")
self.wrist_img = cv2.imread("/home/fishros/hdx/tool/dataset_ori/imgs/0_0.jpg")
# 初始关节角度
self.joint_angles = np.array([244.021978021978,39.32967032967033,17.208791208791208,291.4725274725275,119.56043956043956,75.05494505494505,0], dtype=np.float32)
time.sleep(0.1) # 避免过于频繁的更新
elif self.state == RobotState.ACTING:
iteration += 1
message = self.create_message()
try:
print(f"Sending observation #{iteration}...")
self.socket.send(message)
# 接收动作响应
try:
action_bytes = self.socket.recv()
except zmq.Again:
print("Receive timeout, retrying...")
continue
actions = self.parse_actions(action_bytes)
print("\nReceived actions:")
for i, action in enumerate(actions):
joint_angles = action[:6]
if action[6] < 0.95:
gripper = "OPEN"
else:
gripper = "CLOSE"
self.state = RobotState.FINISHED # 修正:使用赋值操作符=
print(f"Action {i+1}: Joints={joint_angles}, Gripper={gripper}")
# 更新关节角度 (模拟运动)
self.joint_angles[:6] += np.random.uniform(-5, 5, 6)
self.joint_angles[:6] = np.clip(self.joint_angles[:6], -180, 180)
# 每8次迭代切换一次夹爪状态
if iteration % 8 == 0:
self.joint_angles[6] = 1.0 - self.joint_angles[6]
except Exception as e:
print(f"Error occurred: {e}")
self.running = False
elif self.state == RobotState.FINISHED:
# 任务完成状态,显示消息但不发送数据
print("Task finished. Press 'i' to initialize or 'a' to start again.")
time.sleep(0.5) # 避免过于频繁的更新
elif self.state == RobotState.IDLE:
# 空闲状态,显示消息但不发送数据
print("Idle state. Press 'a' to start sending messages.")
time.sleep(0.5) # 避免过于频繁的更新
cv2.destroyAllWindows()
self.socket.close()
self.context.term()
if __name__ == "__main__":
client = Client()
client.run()
print("Client shutdown complete.")接受数据时报错(hdx) (hdx) fishros@fishros-linux:~/hdx/DexGraspVLA/agent$ python hdx_obs_output_socket.py
Client connected to port 15558
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
Idle state. Press 'a' to start sending messages.
State changed to ACTING
Sending observation #1...
Received actions:
Error occurred: 'NoneType' object is not iterable
Client shutdown complete.