openvla / run3.py
stellanwu's picture
Upload folder using huggingface_hub
1975a6f verified
Raw
History Blame Contribute Delete
6.28 kB
import pyrealsense2 as rs
import numpy as np
import cv2
import os
import time
import paramiko
from datetime import datetime
import socket
import urx
import math
# 服务器配置 - 请根据实际情况修改
SERVER_CONFIG = {
'host': '10.103.69.239',
'port': 27277,
'username': 'chubin',
'password': 'ZQSg7YOAeCqGPRky=Aa0#',
'remote_dir': '/mnt/disk_3/chubin/img/'
}
REMOTE_HOST = "localhost" # 控制命令远程主机
REMOTE_PORT = 9999 # 控制端口
UR_ROBOT_IP = "192.168.10.99" # 机械臂IP
instruction = input("请输入任务描述:").strip()
# 配置摄像头管道
def configure_camera():
pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
profile = pipeline.start(config)
depth_sensor = profile.get_device().first_depth_sensor()
depth_scale = depth_sensor.get_depth_scale()
print(f"深度比例系数为: {depth_scale}")
align_to = rs.stream.color
align = rs.align(align_to)
return pipeline, align
# 创建视频写入器
def create_video_writer(width, height, fps=30):
if not os.path.exists("videos"):
os.makedirs("videos")
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
color_filename = f"videos/color_{timestamp}.avi"
depth_filename = f"videos/depth_{timestamp}.avi"
fourcc = cv2.VideoWriter_fourcc(*'XVID')
color_out = cv2.VideoWriter(color_filename, fourcc, fps, (width, height))
depth_out = cv2.VideoWriter(depth_filename, fourcc, fps, (width, height))
return color_out, depth_out
# 上传文件到服务器
def upload_to_server(local_path, remote_path, server_config):
try:
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(
server_config['host'],
port=server_config['port'],
username=server_config['username'],
password=server_config['password']
)
sftp = ssh.open_sftp()
sftp.put(local_path, remote_path)
print(f"文件 {local_path} 已上传到 {remote_path}")
sftp.close()
ssh.close()
return True
except Exception as e:
print(f"上传失败: {e}")
return False
# 获取机械臂控制动作
def get_action_from_remote(instruction: str) -> np.ndarray:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.connect((REMOTE_HOST, REMOTE_PORT))
sock.sendall(instruction.encode())
data = sock.recv(4096)
action_str = data.decode()
action = np.array([float(x) for x in action_str.split(",") if x.strip() != ""])
return action[:-1]
# 机械臂控制
def move_robot(action, ur):
x, y, z, rx, ry, rz = action
#x = x * 1000 # 转换为毫米
#y = y * 1000
#z = (-0.4 + z) * 1000
rx = 2.578 + 1.752 + rx
ry = -0.771 + ry
rz = 4.48 + rz
ur.movej((x, y, z, rx, ry, rz), 0.5, 0.05, relative=False, wait=False)
print("机械臂动作已执行:", action)
# 主程序
def main():
# 连接到机械臂
ur = urx.Robot(UR_ROBOT_IP)
# 配置摄像头和视频写入器
pipeline, align = configure_camera()
color_out, depth_out = create_video_writer(640, 480)
print("正在等待控制指令... 按 'q' 退出")
recording = True
last_capture_time = 0
capture_interval = 5 # 每2秒截图一次
while True:
frames = pipeline.wait_for_frames()
aligned_frames = align.process(frames)
aligned_depth_frame = aligned_frames.get_depth_frame()
color_frame = aligned_frames.get_color_frame()
if not aligned_depth_frame or not color_frame:
continue
depth_image = np.asanyarray(aligned_depth_frame.get_data())
color_image = np.asanyarray(color_frame.get_data())
# 显示录制状态
if recording:
cv2.putText(color_image, "REC", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
color_out.write(color_image)
depth_out.write(cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET))
# 每隔一段时间截图并上传
current_time = time.time()
if current_time - last_capture_time >= capture_interval:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
color_capture_path = f"captures/color_capture_{timestamp}.jpg"
depth_capture_path = f"captures/depth_capture_{timestamp}.jpg"
# 保存截图
cv2.imwrite(color_capture_path, color_image)
cv2.imwrite(depth_capture_path, depth_image)
# 上传文件到服务器
remote_color_path = os.path.join(SERVER_CONFIG['remote_dir'], f"color_capture_{timestamp}.jpg")
remote_depth_path = os.path.join(SERVER_CONFIG['remote_dir'], f"depth_capture_{timestamp}.jpg")
upload_to_server(color_capture_path, remote_color_path, SERVER_CONFIG)
upload_to_server(depth_capture_path, remote_depth_path, SERVER_CONFIG)
last_capture_time = current_time
# 控制机械臂
if instruction:
print(f"发送任务到远程模型服务器...")
action = get_action_from_remote(instruction)
print(f"接收到动作: {action}")
move_robot(action, ur)
# 退出控制
key = cv2.waitKey(1)
if key & 0xFF == ord('q'):
break
elif key & 0xFF == ord('s'): # 按s暂停/继续录制
recording = not recording
print("录制已暂停" if not recording else "录制继续")
# 释放资源
cv2.destroyAllWindows()
pipeline.stop()
color_out.release()
depth_out.release()
ur.close()
if __name__ == "__main__":
main()