cmvr-es-cli/cmvr/biohead_client.py

import grpc
import time
from typing import Tuple
from datetime import datetime

from .enums import CMVRErrorCode
from .models import FacialExpressionState

class BioHeadClient:
    """仿生头客户端"""

    def __init__(self, device_id: str, stub):
        self.device_id = device_id
        self.stub = stub
        self.generated = None  # 推迟导入
        self.request_queue = []  # 请求队列
        self.response_iterator = None  # 响应迭代器
        self.stream_active = False  # 流是否活跃

    def _import_generated(self):
        """推迟导入 generated"""
        if self.generated is None:
            try:
                # 从正确的路径导入生成的模块
                from generated.cmvr.api import common_pb2, biohead_command_pb2
                self.generated = type('GeneratedModules', (), {
                    'common_pb2': common_pb2,
                    'biohead_command_pb2': biohead_command_pb2
                })
            except ImportError as e:
                print(f"导入生成的模块失败: {e}")
                print("请确保已生成 protobuf 代码")
                raise
        return self.generated

    def _create_command_header(self):
        """创建命令头"""
        generated = self._import_generated()
        header = generated.common_pb2.CommandHeader.Request()
        header.device_id = self.device_id
        # 使用正确的方式设置时间戳
        from google.protobuf.timestamp_pb2 import Timestamp
        timestamp = Timestamp()
        timestamp.GetCurrentTime()
        header.timestamp.CopyFrom(timestamp)
        return header

    def set_expression(self, expression: FacialExpressionState) -> CMVRErrorCode:
        """设置表情"""
        try:
            generated = self._import_generated()
            # 创建请求
            request = generated.biohead_command_pb2.SetFacialExpression.Request()
            request.header.CopyFrom(self._create_command_header())

            # 填充表情数据
            expr = request.expression
            expr.eyebrow.left_outside_y = expression.left_eyebrow_outside_y
            expr.eyebrow.left_inside_y = expression.left_eyebrow_inside_y
            expr.eyebrow.right_outside_y = expression.right_eyebrow_outside_y
            expr.eyebrow.right_inside_y = expression.right_eyebrow_inside_y

            expr.eyelid.left_upper_y = expression.left_eye_upper_lid_y
            expr.eyelid.left_lower_y = expression.left_eye_lower_lid_y
            expr.eyelid.right_upper_y = expression.right_eye_upper_lid_y
            expr.eyelid.right_lower_y = expression.right_eye_lower_lid_y

            expr.eyeball.left_x = expression.left_eye_ball_x
            expr.eyeball.left_y = expression.left_eye_ball_y
            expr.eyeball.right_x = expression.right_eye_ball_x
            expr.eyeball.right_y = expression.right_eye_ball_y

            expr.nose.left_y = expression.left_nose_y
            expr.nose.right_y = expression.right_nose_y

            expr.mouth.upper_lip_y = expression.upper_lip_y
            expr.mouth.upper_lip_z = expression.upper_lip_z
            expr.mouth.lower_lip_y = expression.lower_lip_y
            expr.mouth.lower_lip_z = expression.lower_lip_z

            expr.mouth.left_lip.upper_x = expression.upper_left_lip_x
            expr.mouth.left_lip.upper_y = expression.upper_left_lip_y
            expr.mouth.left_lip.corner_x = expression.left_corner_lip_x
            expr.mouth.left_lip.corner_y = expression.left_corner_lip_y
            expr.mouth.left_lip.lower_x = expression.lower_left_lip_x
            expr.mouth.left_lip.lower_y = expression.lower_left_lip_y

            expr.mouth.right_lip.upper_x = expression.upper_right_lip_x
            expr.mouth.right_lip.upper_y = expression.upper_right_lip_y
            expr.mouth.right_lip.corner_x = expression.right_corner_lip_x
            expr.mouth.right_lip.corner_y = expression.right_corner_lip_y
            expr.mouth.right_lip.lower_x = expression.lower_right_lip_x
            expr.mouth.right_lip.lower_y = expression.lower_right_lip_y

            expr.jaw.x = expression.jaw_x
            expr.jaw.y = expression.jaw_y

            # 发送请求
            response = self.stub.SetExpression(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"设置表情失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def start_stream(self):
        """开始流式会话"""
        try:
            if not self.stream_active:
                # 创建一个请求迭代器生成器
                self.request_queue = []
                self.stream_active = True
                self.response_iterator = self.stub.StreamExpression(self._stream_request_generator())
            return CMVRErrorCode.CMVR_SUCCESS
        except grpc.RpcError as e:
            print(f"开始流式会话失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def _stream_request_generator(self):
        """流式请求生成器"""
        generated = self._import_generated()
        while self.stream_active:
            if self.request_queue:
                request = self.request_queue.pop(0)
                yield request
                if request.eof:
                    break
            else:
                # 如果没有请求，等待一小段时间
                time.sleep(0.01)
            time.sleep(0.1)

    def stream_expression(self, expression, is_eof=False):
        """流式控制表情"""
        try:
            if not self.stream_active:
                result = self.start_stream()
                if result != CMVRErrorCode.CMVR_SUCCESS:
                    return result

            generated = self._import_generated()

            # 创建请求
            request = generated.biohead_command_pb2.StreamFacialExpression.Request()
            request.header.CopyFrom(self._create_command_header())
            request.expr.CopyFrom(expression)
            request.eof = is_eof

            # 将请求添加到队列
            self.request_queue.append(request)

            # 读取最新的响应
            try:
                feedback = next(self.response_iterator)
                return CMVRErrorCode.CMVR_SUCCESS if feedback.header.success else CMVRErrorCode.CMVR_RPC_FAILED
            except StopIteration:
                return CMVRErrorCode.CMVR_SUCCESS

        except grpc.RpcError as e:
            print(f"流式控制表情失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def end_stream(self):
        """结束流式会话"""
        try:
            if self.stream_active:
                # 发送结束信号
                generated = self._import_generated()
                request = generated.biohead_command_pb2.StreamFacialExpression.Request()
                request.header.CopyFrom(self._create_command_header())
                request.eof = True

                self.request_queue.append(request)

                # 等待所有响应
                for feedback in self.response_iterator:
                    if not feedback.header.success:
                        return CMVRErrorCode.CMVR_RPC_FAILED

                self.stream_active = False
                self.request_queue = []
                self.response_iterator = None

            return CMVRErrorCode.CMVR_SUCCESS
        except grpc.RpcError as e:
            print(f"结束流式会话失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def get_system_status(self):
        """获取系统状态"""
        try:
            generated = self._import_generated()
            request = generated.biohead_command_pb2.GetStatus.Request()
            request.header.CopyFrom(self._create_command_header())

            response = self.stub.GetSystemStatus(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"获取系统状态失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def emergency_stop(self):
        """紧急停止"""
        try:
            generated = self._import_generated()
            request = generated.biohead_command_pb2.EmergencyStop.Request()
            request.header.CopyFrom(self._create_command_header())

            response = self.stub.EmergencyStop(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"紧急停止失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def close(self):
        """关闭连接"""
        if self.stream_active:
            self.end_stream()