179 lines
6.0 KiB
Python
179 lines
6.0 KiB
Python
# coding=utf-8
|
||
import json
|
||
import logging
|
||
import os
|
||
import socket
|
||
import time
|
||
import sys
|
||
import subprocess
|
||
import re
|
||
import numpy as np
|
||
import cv2
|
||
import pyrealsense2 as rs
|
||
|
||
from libs.log_setting import CommonLog
|
||
from libs.auxiliary import create_folder_with_date, get_ip, popup_message
|
||
from save_poses import save_matrices_to_csv
|
||
|
||
# 把 .so 所在目录加入 Python 路径
|
||
sys.path.append("/home/lgv/cmvr/cmvr-es/cmake-build-debug/example")
|
||
from robot_wrapper import Robot
|
||
|
||
# -------------------- 全局配置 --------------------
|
||
TARGET_SERIAL = "243122075614" # 相机序列号
|
||
EXEC_PATH = "/home/lgv/cmvr/cmvr-es/cmake-build-debug/example/solve_fk"
|
||
ROBOT_CONFIG = "/home/lgv/cmvr/cmvr-es/config/cabin_robot.xml"
|
||
|
||
logger = logging.getLogger(__name__)
|
||
logger = CommonLog(logger)
|
||
|
||
robot = Robot(ROBOT_CONFIG, "hc01")
|
||
output_dir = create_folder_with_date() # 存储目录
|
||
count = 1
|
||
|
||
|
||
# -------------------- 工具函数 --------------------
|
||
def get_pose(joint_positions, base_link="PELVIS_S", target_link="R_FINGER_TIP", exec_path=EXEC_PATH):
|
||
"""通过 solve_fk 求解 4x4 位姿矩阵"""
|
||
joint_str = [str(j) for j in joint_positions]
|
||
cmd = [exec_path, base_link, target_link, *joint_str]
|
||
|
||
try:
|
||
res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||
check=True, text=True)
|
||
except subprocess.CalledProcessError as e:
|
||
raise RuntimeError(f"solve_fk 调用失败: {e.stdout}")
|
||
|
||
lines = res.stdout.splitlines()
|
||
start_idx = next((i for i, l in enumerate(lines) if re.search(r"变换矩阵|T\s*:", l)), None)
|
||
if start_idx is None or start_idx + 4 >= len(lines):
|
||
raise RuntimeError("未找到 4x4 变换矩阵:\n" + res.stdout)
|
||
|
||
try:
|
||
mat = np.array([[float(x) for x in lines[start_idx + 1 + r].split()] for r in range(4)], dtype=np.float64)
|
||
if mat.shape != (4, 4):
|
||
raise ValueError
|
||
except Exception:
|
||
raise RuntimeError("矩阵解析失败:\n" + res.stdout)
|
||
return mat
|
||
|
||
|
||
def save_pose_and_image(pose, image, index):
|
||
"""保存位姿矩阵和对应的图片"""
|
||
try:
|
||
# 保存位姿（追加模式 'a'）
|
||
pose_file = os.path.join(output_dir, "RobotToolPose.csv")
|
||
save_matrices_to_csv([pose], pose_file) # 假设 save_matrices_to_csv 支持 mode 参数
|
||
|
||
# 保存图片
|
||
img_file = os.path.join(output_dir, f"{index}.jpg")
|
||
cv2.imwrite(img_file, image)
|
||
|
||
logger.info(f"===采集第{index}次数据！ pose 已保存到 {pose_file}, 图片保存到 {img_file}")
|
||
except Exception as e:
|
||
logger.error(f"保存数据失败: {e}")
|
||
|
||
|
||
def callback_and_save(frame):
|
||
"""显示视频帧，按 's' 保存，按 'q' 返回 False 退出"""
|
||
global count
|
||
cv2.imshow("Capture_Video2", frame)
|
||
scaling_factor = 2.0
|
||
cv_img = cv2.resize(frame, None, fx=scaling_factor, fy=scaling_factor,
|
||
interpolation=cv2.INTER_AREA)
|
||
cv2.imshow("Capture_Video", cv_img)
|
||
|
||
k = cv2.waitKey(30) & 0xFF
|
||
if k == ord('s'):
|
||
try:
|
||
# 获取关节角 & 位姿
|
||
js = robot.getJointQ('right')
|
||
logger.info(f"当前关节角: {js}")
|
||
|
||
pose = get_pose(js, base_link="PELVIS_S", target_link="R_FINGER_TIP")
|
||
logger.info(f"获取位姿成功:\n{pose}")
|
||
logger.info("请按 p 失能机器人，并手动移动机器人到其他位姿")
|
||
|
||
save_pose_and_image(pose, cv_img, count)
|
||
count += 1
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取或保存数据失败: {e}")
|
||
elif k == ord('q'):
|
||
logger.info("检测到 'q'，退出程序")
|
||
return False
|
||
elif k == ord('o'):
|
||
robot.torqueOn()
|
||
logger.info("robot.torqueOn ，请勿触摸机器人")
|
||
elif k == ord('p'):
|
||
logger.info("robot.torqueOff")
|
||
robot.torqueOff()
|
||
|
||
|
||
return True
|
||
|
||
|
||
|
||
|
||
def start_pipeline(serial=TARGET_SERIAL, depth_size=(1280, 720), color_size=(1280, 720), fps: int = 30):
|
||
ctx = rs.context()
|
||
devices = ctx.query_devices()
|
||
if len(devices) == 0:
|
||
raise RuntimeError("未检测到任何 RealSense 相机")
|
||
|
||
if not any(dev.get_info(rs.camera_info.serial_number) == serial for dev in devices):
|
||
raise RuntimeError(f"未找到序列号为 {serial} 的 RealSense 相机")
|
||
|
||
pipeline = rs.pipeline()
|
||
cfg = rs.config()
|
||
cfg.enable_device(serial)
|
||
|
||
cfg.enable_stream(rs.stream.depth, depth_size[0], depth_size[1], rs.format.z16, fps)
|
||
cfg.enable_stream(rs.stream.color, color_size[0], color_size[1], rs.format.bgr8, fps)
|
||
|
||
pipeline_profile = pipeline.start(cfg)
|
||
|
||
# 相机内参
|
||
depth_intr = pipeline_profile.get_stream(rs.stream.depth).as_video_stream_profile().get_intrinsics()
|
||
color_intr = pipeline_profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
|
||
|
||
# 建议对齐到彩色流
|
||
align = rs.align(rs.stream.color)
|
||
# align = rs.align(rs.stream.depth)
|
||
|
||
return pipeline, align, depth_intr, color_intr
|
||
def collect_data():
|
||
"""循环采集机器人位姿与相机图像"""
|
||
logger.info("开始机器人数据采集程序，版本 V1.0.0")
|
||
|
||
pipeline, align, depth_intr, color_intr = start_pipeline()
|
||
global count
|
||
count = 1
|
||
|
||
try:
|
||
while True:
|
||
frames = pipeline.wait_for_frames()
|
||
aligned_frames = align.process(frames)
|
||
color_frame = aligned_frames.get_color_frame()
|
||
depth_frame = aligned_frames.get_depth_frame()
|
||
if not color_frame or not depth_frame:
|
||
continue
|
||
|
||
# 转换为 numpy
|
||
color_image = np.asanyarray(color_frame.get_data())
|
||
depth_image = np.asanyarray(depth_frame.get_data())
|
||
|
||
# 调用一次 callback_and_save
|
||
if not callback_and_save(color_image):
|
||
break
|
||
|
||
finally:
|
||
pipeline.stop()
|
||
cv2.destroyAllWindows()
|
||
logger.info("相机已关闭，窗口已销毁")
|
||
|
||
|
||
# -------------------- 主入口 --------------------
|
||
if __name__ == '__main__':
|
||
collect_data()
|