cmvr-es/python/hand_eye_calibration/collect_data.py

# 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 pathlib import Path

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"
REPO_ROOT = Path(__file__).resolve().parents[2]
ROBOT_CONFIG = str(REPO_ROOT / "cmvr-es/config/cmvr_es.pb.txt")

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()