kuavo-humanoid-sdk-ws 1.2.2__20250922181225-py3-none-any.whl

kuavo_humanoid_sdk/kuavo/core/ros/vision.py

@@ -0,0 +1,280 @@
+#! /usr/bin/env python3
+# coding: utf-8
+
+import copy
+import time
+import numpy as np
+from collections import deque
+from typing import Tuple, Optional
+from transforms3d import quaternions, euler
+
+from kuavo_humanoid_sdk.kuavo.core.ros.param import make_robot_param, EndEffectorType
+from kuavo_humanoid_sdk.common.logger import SDKLogger
+from kuavo_humanoid_sdk.interfaces.data_types import (AprilTagData, AprilTagDetection)
+import roslibpy
+from kuavo_humanoid_sdk.common.websocket_kuavo_sdk import WebSocketKuavoSDK
+
+class KuavoRobotVisionCoreWebsocket:
+    """Handles vision-related data processing for Kuavo humanoid robot.
+    
+    Attributes:
+        tf_client (roslibpy.Topic): TF client for coordinate transformations
+        tf_publisher (roslibpy.Topic): TF publisher for broadcasting transforms
+    """
+    
+    def __init__(self):
+        """Initializes vision system components including TF and AprilTag subscribers."""
+        if not hasattr(self, '_initialized'):
+            # Initialize WebSocket connection
+            self.websocket = WebSocketKuavoSDK()
+            
+            # Initialize AprilTag subscribers
+            self._apriltag_data_camera_sub = roslibpy.Topic(
+                self.websocket.client,
+                '/tag_detections',
+                'apriltag_ros/AprilTagDetectionArray'
+            )
+            self._apriltag_data_camera_sub.subscribe(self._apriltag_data_callback_camera)
+            
+            self._apriltag_data_base_sub = roslibpy.Topic(
+                self.websocket.client,
+                '/robot_tag_info',
+                'apriltag_ros/AprilTagDetectionArray'
+            )
+            self._apriltag_data_base_sub.subscribe(self._apriltag_data_callback_base)
+
+            self._apriltag_data_odom_sub = roslibpy.Topic(
+                self.websocket.client,
+                '/robot_tag_info_odom',
+                'apriltag_ros/AprilTagDetectionArray'
+            )
+            self._apriltag_data_odom_sub.subscribe(self._apriltag_data_callback_odom)
+
+            # Initialize data structures
+            self._apriltag_data_from_camera = AprilTagData(
+                id = [],
+                size = [],
+                pose = []
+            )
+
+            self._apriltag_data_from_base = AprilTagData(
+                id = [],
+                size = [],
+                pose = []
+            )
+
+            self._apriltag_data_from_odom = AprilTagData(
+                id = [],
+                size = [],
+                pose = []
+            )
+            
+            self._initialized = True
+
+    def _tf_callback(self, msg):
+        """Callback for TF messages.
+        
+        Args:
+            msg: TF message containing transforms
+        """
+        for transform in msg['transforms']:
+            key = (transform['header']['frame_id'], transform['child_frame_id'])
+            self._transforms[key] = transform
+
+    def _apriltag_data_callback_camera(self, data):
+        """Callback for processing AprilTag detections from camera.
+        
+        Args:
+            data (dict): Raw detection data from camera
+        """
+        # 清空之前的数据
+        self._apriltag_data_from_camera.id = []
+        self._apriltag_data_from_camera.size = []
+        self._apriltag_data_from_camera.pose = []
+        
+        # 处理每个标签检测
+        for detection in data['detections']:
+            # 添加ID
+            for id in detection['id']:
+                self._apriltag_data_from_camera.id.append(id)
+            
+            # 添加尺寸
+            if detection.get('size') and len(detection['size']) > 0:
+                self._apriltag_data_from_camera.size.append(detection['size'][0])
+            else:
+                self._apriltag_data_from_camera.size.append(0.0)
+            
+            
+            # Convert pose dict to AprilTagDetection
+            pose_dict = detection['pose']['pose']['pose']
+
+            tag_detection = AprilTagDetection(
+                position=AprilTagDetection.Point(
+                    x=float(pose_dict['position']['x']), 
+                    y=float(pose_dict['position']['y']), 
+                    z=float(pose_dict['position']['z'])
+                    ), 
+                orientation=AprilTagDetection.Quaternion(
+                    x=float(pose_dict['orientation']['x']), 
+                    y=float(pose_dict['orientation']['y']), 
+                    z=float(pose_dict['orientation']['z']), 
+                    w=float(pose_dict['orientation']['w'])
+                ))
+            # 添加姿态
+            self._apriltag_data_from_camera.pose.append(tag_detection)
+            
+    def _apriltag_data_callback_base(self, data):
+        """Callback for processing AprilTag detections from base link.
+        
+        Args:
+            data (dict): Raw detection data from base frame
+        """
+        # 清空之前的数据
+        self._apriltag_data_from_base.id = []
+        self._apriltag_data_from_base.size = []
+        self._apriltag_data_from_base.pose = []
+        
+        # 处理每个标签检测
+        for detection in data['detections']:
+            # 添加ID
+            for id in detection['id']:
+                self._apriltag_data_from_base.id.append(id)
+            
+            # 添加尺寸
+            if detection.get('size') and len(detection['size']) > 0:
+                self._apriltag_data_from_base.size.append(detection['size'][0])
+            else:
+                self._apriltag_data_from_base.size.append(0.0)
+            
+                        # Convert pose dict to AprilTagDetection
+            pose_dict = detection['pose']['pose']['pose']
+            
+            tag_detection = AprilTagDetection(
+                position=AprilTagDetection.Point(
+                    x=float(pose_dict['position']['x']), 
+                    y=float(pose_dict['position']['y']), 
+                    z=float(pose_dict['position']['z'])
+                    ), 
+                orientation=AprilTagDetection.Quaternion(
+                    x=float(pose_dict['orientation']['x']), 
+                    y=float(pose_dict['orientation']['y']), 
+                    z=float(pose_dict['orientation']['z']), 
+                    w=float(pose_dict['orientation']['w'])
+                ))
+            
+            # 添加姿态
+            self._apriltag_data_from_base.pose.append(tag_detection)
+ 
+    def _apriltag_data_callback_odom(self, data):
+        """Callback for processing AprilTag detections from odom frame.
+        
+        Args:
+            data (dict): Raw detection data from odom frame
+        """
+        # 清空之前的数据
+        self._apriltag_data_from_odom.id = []
+        self._apriltag_data_from_odom.size = []
+        self._apriltag_data_from_odom.pose = []
+        
+        # 处理每个标签检测
+        for detection in data['detections']:
+            # 添加ID
+            for id in detection['id']:
+                self._apriltag_data_from_odom.id.append(id)
+            
+            # 添加尺寸
+            if detection.get('size') and len(detection['size']) > 0:
+                self._apriltag_data_from_odom.size.append(detection['size'][0])
+            else:
+                self._apriltag_data_from_odom.size.append(0.0)
+            
+            # Convert pose dict to AprilTagDetection
+            pose_dict = detection['pose']['pose']['pose']
+            
+            tag_detection = AprilTagDetection(
+                position=AprilTagDetection.Point(
+                    x=float(pose_dict['position']['x']), 
+                    y=float(pose_dict['position']['y']), 
+                    z=float(pose_dict['position']['z'])
+                    ), 
+                orientation=AprilTagDetection.Quaternion(
+                    x=float(pose_dict['orientation']['x']), 
+                    y=float(pose_dict['orientation']['y']), 
+                    z=float(pose_dict['orientation']['z']), 
+                    w=float(pose_dict['orientation']['w'])
+                ))
+            
+            # 添加姿态
+            self._apriltag_data_from_odom.pose.append(tag_detection)
+
+    @property 
+    def apriltag_data_from_camera(self) -> AprilTagData:
+        """AprilTag detection data in camera coordinate frame.
+        
+        Returns:
+            AprilTagData: Contains lists of tag IDs, sizes and poses
+        """
+        return self._apriltag_data_from_camera
+    
+    @property
+    def apriltag_data_from_base(self) -> AprilTagData:
+        """AprilTag detection data in base_link coordinate frame.
+        
+        Returns:
+            AprilTagData: Contains lists of tag IDs, sizes and poses
+        """
+        return self._apriltag_data_from_base
+    
+    @property
+    def apriltag_data_from_odom(self) -> AprilTagData:
+        """AprilTag detection data in odom coordinate frame.
+        
+        Returns:
+            AprilTagData: Contains lists of tag IDs, sizes and transformed poses
+        """
+        return self._apriltag_data_from_odom
+
+    def _get_data_by_id(self, target_id: int, data_source: str = "base") -> Optional[dict]:
+        """Retrieves AprilTag data by specific ID from selected source.
+        
+        Args:
+            target_id (int): AprilTag ID to search for
+            data_source (str): Data source selector, valid options: 
+                "camera", "base", "odom"
+        
+        Returns:
+            Optional[dict]: Dictionary containing 'sizes' and 'poses' lists if found,
+                None if no matching data
+        """
+        data_map = {
+            "camera": self._apriltag_data_from_camera,
+            "base": self._apriltag_data_from_base,
+            "odom": self._apriltag_data_from_odom
+        }
+        
+        if data_source not in data_map:
+            raise ValueError(f"Invalid data source: {data_source}")
+            
+        data = data_map[data_source]
+        
+        # 查找匹配的ID
+        try:
+            idx = data.id.index(target_id)
+            return {
+                'sizes': [data.size[idx]],
+                'poses': [data.pose[idx]]
+            }
+        except ValueError:
+            return None
+
+# if __name__ == "__main__":
+
+#     kuavo_robot_vision_core = KuavoRobotVisionCoreWebsocket()
+#     time.sleep(5)
+#     print("apriltag_data_from_camera:")
+#     print(kuavo_robot_vision_core.apriltag_data_from_camera)
+#     print("apriltag_data_from_base:")
+#     print(kuavo_robot_vision_core.apriltag_data_from_base)
+#     print("apriltag_data_from_odom:")
+#     print(kuavo_robot_vision_core.apriltag_data_from_odom)
+#     rospy.spin()

kuavo_humanoid_sdk/kuavo/core/ros_env.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python3
+# coding: utf-8
+
+import os
+import subprocess
+import atexit
+from kuavo_humanoid_sdk.common.logger import SDKLogger
+from kuavo_humanoid_sdk.common.websocket_kuavo_sdk import WebSocketKuavoSDK
+import roslibpy
+
+
+class KuavoROSEnvWebsocket:
+    _instance = None
+    _processes = []  # Store all subprocess instances
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(KuavoROSEnvWebsocket, cls).__new__(cls)
+            cls._instance._initialized = False
+            cls._instance._init_success = False  # Add flag to track successful Init() call
+            # Register cleanup handler
+            atexit.register(cls._cleanup_processes)
+        return cls._instance
+
+    def __init__(self):
+        if not self._initialized:
+            self._initialized = True
+            self.websocket = None
+
+    @classmethod
+    def _cleanup_processes(cls):
+        """Cleanup all registered processes on exit"""
+        for process in cls._processes:
+            if process.poll() is None:  # Process is still running
+                process.terminate()
+                try:
+                    process.wait(timeout=3)  # Wait for process to terminate
+                except subprocess.TimeoutExpired:
+                    process.kill()  # Force kill if not terminated
+        cls._processes.clear()
+
+    def _get_kuavo_ws_root(self) -> str:
+        # For WebSocket version, we'll use environment variable instead of ROS param
+        model_path = os.environ.get('KUAVO_MODEL_PATH')
+        if not model_path:
+            raise Exception("KUAVO_MODEL_PATH environment variable not found")
+        
+        if not os.path.exists(model_path):
+            raise Exception(f"Model path {model_path} not found")
+            
+        # ws 
+        return model_path.replace('/src/kuavo_assets/models', '')
+    
+    def Init(self) -> bool:
+        """
+        Initialize the WebSocket environment.
+        Raises:
+            Exception: If the KUAVO_MODEL_PATH environment variable is not found or the path is not valid.
+        """
+        # if generate docs, skip init.
+        if 'GEN_KUAVO_HUMANOID_SDK_DOCS' in os.environ:
+            return True
+
+        # Return directly if already initialized successfully
+        if self._init_success:
+            return True
+        
+        # Check if WebSocket server is running
+        try:
+            
+            self.websocket = WebSocketKuavoSDK()
+            if not self.websocket.client.is_connected:
+                print(f"\033[31m\nError: Can't connect to WebSocket server. Please ensure the server is running.\033[0m"
+                      "\nMaybe manually launch the app first?"
+                      "\n - for example(sim): roslaunch humanoid_controller load_kuavo_mujoco_sim.launch, "
+                      "\n - for example(real): roslaunch humanoid_controller load_kuavo_real.launch\n")
+                exit(1)
+        except Exception as e:
+            print(f"\033[31m\nError: Failed to connect to WebSocket server: {e}\033[0m")
+            exit(1)
+        
+        # Only check nodes exist when Init SDK, if not, tips user manually launch nodes.
+        deps_nodes = ['/humanoid_gait_switch_by_name']
+        for node in deps_nodes:
+            if not self.check_rosnode_exists(node):
+                print(f"\033[31m\nError: Node {node} not found. Please launch it manually.\033[0m")
+                exit(1)
+                
+        self._init_success = True  # Set flag after successful initialization
+
+        return True
+    
+    def _launch_ros_node(self, node_name, launch_cmd, log_name):
+        """Launch a ROS node with the given command and log the output.
+        
+        Args:
+            node_name (str): Name of the node to launch
+            launch_cmd (str): Full launch command including source and roslaunch
+            log_name (str): Name for the log file
+            
+        Raises:
+            Exception: If node launch fails
+        """
+        # Launch in background and check if successful
+        try:
+            os.makedirs('/var/log/kuavo_humanoid_sdk', exist_ok=True)
+            log_path = f'/var/log/kuavo_humanoid_sdk/{log_name}.log'
+        except PermissionError:
+            os.makedirs('log/kuavo_humanoid_sdk', exist_ok=True)
+            log_path = f'log/kuavo_humanoid_sdk/{log_name}.log'
+            
+        with open(log_path, 'w') as log_file:
+            process = subprocess.Popen(launch_cmd, shell=True, executable='/bin/bash', stdout=log_file, stderr=log_file)
+            self._processes.append(process)  # Add process to tracking list
+            
+        if process.returncode is not None and process.returncode != 0:
+            raise Exception(f"Failed to launch {node_name}, return code: {process.returncode}")
+
+        SDKLogger.info(f"{node_name} launched successfully")
+
+    def _get_setup_file(self, ws_root=None):
+        """Get the appropriate ROS setup file path based on shell type.
+        
+        Args:
+            ws_root (str, optional): ROS workspace root path. If None, uses ROS_WORKSPACE.
+            
+        Returns:
+            str: Path to the setup file
+            
+        Raises:
+            Exception: If setup file not found
+        """
+        is_zsh = 'zsh' in os.environ.get('SHELL', '')
+        
+        if ws_root is None:
+            ws_root = os.environ['ROS_WORKSPACE']
+            
+        setup_files = {
+            'zsh': os.path.join(ws_root, 'devel/setup.zsh'),
+            'bash': os.path.join(ws_root, 'devel/setup.bash')
+        }
+        
+        setup_file = setup_files['zsh'] if is_zsh else setup_files['bash']
+        if not os.path.exists(setup_file):
+            setup_file = setup_file.replace('devel', 'install')
+            if not os.path.exists(setup_file):
+                raise Exception(f"Setup file not found in either devel or install: {setup_file}")
+                
+        return setup_file
+
+    def launch_ik_node(self):
+        # nodes: /arms_ik_node
+        # services: /ik/two_arm_hand_pose_cmd_srv, /ik/fk_srv
+        try:
+            if not self.websocket or not self.websocket.client.is_connected:
+                raise Exception("WebSocket server is not running")
+        except Exception as e:
+            raise Exception(f"WebSocket server is not running: {e}")
+
+        # Check if IK node and services exist
+        try:
+            # Check if arms_ik_node is running using roslibpy
+            service = roslibpy.Service(self.websocket.client, '/rosnode/list', 'rosapi/Nodes')
+            response = service.call({})
+            nodes = response.get('nodes', [])
+            
+            if '/arms_ik_node' not in nodes:
+                # Launch IK node if not running
+                kuavo_ws_root = self._get_kuavo_ws_root()
+                setup_file = self._get_setup_file(kuavo_ws_root)
+                source_cmd = f"source {setup_file}"
+                
+                # Get robot version from environment variable
+                robot_version = os.environ.get('ROBOT_VERSION')
+                if robot_version is None:
+                    raise Exception("Failed to get ROBOT_VERSION from environment variables")
+                
+                # Launch IK node
+                launch_cmd = f"roslaunch motion_capture_ik ik_node.launch robot_version:={robot_version}"
+                full_cmd = f"{source_cmd} && {launch_cmd}"
+                
+                self._launch_ros_node('IK node', full_cmd, 'launch_ik')
+
+            return True
+
+        except Exception as e:
+            raise Exception(f"Failed to verify IK node and services: {e}")
+        
+
+    def launch_gait_switch_node(self)-> bool:
+        """Verify that the gait switch node is running, launch if not."""
+        try:
+            # Check if node exists using roslibpy
+            service = roslibpy.Service(self.websocket.client, '/rosnode/list', 'rosapi/Nodes')
+            response = service.call({})
+            nodes = response.get('nodes', [])
+            
+            if '/humanoid_gait_switch_by_name' not in nodes:
+                kuavo_ws_root = self._get_kuavo_ws_root()
+                setup_file = self._get_setup_file(kuavo_ws_root)
+                source_cmd = f"source {setup_file}"
+    
+                # Launch gait switch node
+                launch_cmd = "roslaunch humanoid_interface_ros humanoid_switch_gait.launch"
+                full_cmd = f"{source_cmd} && {launch_cmd}"
+                
+                self._launch_ros_node('Gait switch node', full_cmd, 'launch_gait_switch')
+
+        except Exception as e:
+            raise Exception(f"Failed to launch gait_switch node: {e}")
+        
+    @staticmethod
+    def check_rosnode_exists(node_name):
+        """Check if a ROS node is running using roslibpy.
+        
+        Args:
+            node_name (str): Name of the node to check
+            
+        Returns:
+            bool: True if node is running, False otherwise
+        """
+        websocket = WebSocketKuavoSDK()
+        try:
+            if not websocket or not websocket.client.is_connected:
+                return False
+                
+            service = roslibpy.Service(websocket.client, '/rosapi/nodes', 'rosapi/Nodes')
+            response = service.call({})
+            nodes = response.get('nodes', [])
+            return node_name in nodes
+        except Exception as e:
+            SDKLogger.error(f"Error checking if node {node_name} exists: {e}")
+            return False