dsh/dual_camera_system/ptz_camera.py

"""
球机(PTZ)控制模块
负责PTZ控制和精确定位
"""

import math
import time
import threading
from typing import Optional, Tuple, Dict
from dataclasses import dataclass

from config import PTZ_CAMERA, PTZ_CONFIG
from dahua_sdk import DahuaSDK, PTZCommand


@dataclass
class PTZPosition:
    """PTZ位置信息"""
    pan: float      # 水平角度 (0-360度)
    tilt: float     # 垂直角度 (-90到90度)
    zoom: float     # 变倍 (1-最大倍数)


class PTZCamera:
    """球机控制类"""
    
    def __init__(self, sdk: DahuaSDK, camera_config: Dict = None):
        """
        初始化球机
        Args:
            sdk: 大华SDK实例
            camera_config: 摄像头配置
        """
        self.sdk = sdk
        self.config = camera_config or PTZ_CAMERA
        self.ptz_config = PTZ_CONFIG
        
        self.login_handle = None
        self.connected = False
        
        # 当前位置
        self.current_position = PTZPosition(pan=0, tilt=0, zoom=1)
        self.position_lock = threading.Lock()
    
    def connect(self) -> bool:
        """
        连接球机
        Returns:
            是否成功
        """
        login_handle, error = self.sdk.login(
            self.config['ip'],
            self.config['port'],
            self.config['username'],
            self.config['password']
        )
        
        if login_handle is None:
            print(f"连接球机失败: IP={self.config['ip']}, 错误码={error}")
            return False
        
        self.login_handle = login_handle
        self.connected = True
        print(f"成功连接球机: {self.config['ip']}")
        return True
    
    def disconnect(self):
        """断开连接"""
        if self.login_handle:
            self.sdk.logout(self.login_handle)
            self.login_handle = None
        self.connected = False
    
    def ptz_control(self, command: int, param1: int = 0, param2: int = 0, 
                    param3: int = 0, stop: bool = False) -> bool:
        """
        PTZ控制
        Args:
            command: 控制命令
            param1-3: 参数
            stop: 是否停止
        Returns:
            是否成功
        """
        if not self.connected:
            return False
        
        return self.sdk.ptz_control(
            self.login_handle,
            self.config['channel'],
            command, param1, param2, param3, stop
        )
    
    def move_up(self, speed: int = 4, stop: bool = False) -> bool:
        """向上移动"""
        return self.ptz_control(PTZCommand.UP, 0, speed, 0, stop)
    
    def move_down(self, speed: int = 4, stop: bool = False) -> bool:
        """向下移动"""
        return self.ptz_control(PTZCommand.DOWN, 0, speed, 0, stop)
    
    def move_left(self, speed: int = 4, stop: bool = False) -> bool:
        """向左移动"""
        return self.ptz_control(PTZCommand.LEFT, 0, speed, 0, stop)
    
    def move_right(self, speed: int = 4, stop: bool = False) -> bool:
        """向右移动"""
        return self.ptz_control(PTZCommand.RIGHT, 0, speed, 0, stop)
    
    def zoom_in(self, speed: int = 4, stop: bool = False) -> bool:
        """放大"""
        return self.ptz_control(PTZCommand.ZOOM_ADD, 0, speed, 0, stop)
    
    def zoom_out(self, speed: int = 4, stop: bool = False) -> bool:
        """缩小"""
        return self.ptz_control(PTZCommand.ZOOM_DEC, 0, speed, 0, stop)
    
    def stop_move(self) -> bool:
        """停止移动"""
        # 发送停止命令
        self.move_up(0, True)
        self.move_left(0, True)
        self.zoom_in(0, True)
        return True
    
    def goto_exact_position(self, pan: float, tilt: float, zoom: int) -> bool:
        """
        三维精确定位
        Args:
            pan: 水平角度 (0-360度)
            tilt: 垂直角度 (-90到90度)
            zoom: 变倍 (1-128)
        Returns:
            是否成功
        """
        # SDK参数: pan 0~3600 (0.1度单位), tilt -1800~1800, zoom 1~128
        param1 = int(pan * 10)      # 转换为0.1度单位
        param2 = int(tilt * 10)     # 转换为0.1度单位
        param3 = int(min(zoom, 128))  # 限制最大128倍
        
        result = self.ptz_control(PTZCommand.EXACTGOTO, param1, param2, param3)
        
        if result:
            with self.position_lock:
                self.current_position = PTZPosition(pan=pan, tilt=tilt, zoom=zoom)
        
        return result
    
    def goto_preset(self, preset_id: int) -> bool:
        """
        转到预置点
        Args:
            preset_id: 预置点ID
        Returns:
            是否成功
        """
        return self.ptz_control(PTZCommand.POINT_GO, 0, preset_id, 0)
    
    def set_preset(self, preset_id: int) -> bool:
        """
        设置预置点
        Args:
            preset_id: 预置点ID
        Returns:
            是否成功
        """
        return self.ptz_control(PTZCommand.POINT_SET, 0, preset_id, 0)
    
    def clear_preset(self, preset_id: int) -> bool:
        """
        清除预置点
        Args:
            preset_id: 预置点ID
        Returns:
            是否成功
        """
        return self.ptz_control(PTZCommand.POINT_CLEAR, 0, preset_id, 0)
    
    def calculate_ptz_position(self, x_ratio: float, y_ratio: float,
                                zoom: int = None) -> Tuple[float, float, int]:
        """
        根据全景画面中的位置计算PTZ角度
        Args:
            x_ratio: X方向比例 (0-1)
            y_ratio: Y方向比例 (0-1)
            zoom: 变倍 (默认使用配置值)
        Returns:
            (pan, tilt, zoom) PTZ位置
        """
        if zoom is None:
            zoom = self.ptz_config['default_zoom']
        
        # 应用坐标偏移校准
        offset_x, offset_y = self.ptz_config['coordinate_offset']
        x_ratio = max(0, min(1, x_ratio + offset_x))
        y_ratio = max(0, min(1, y_ratio + offset_y))
        
        # 从配置获取视野参数
        pan_range = self.ptz_config.get('pan_range', (0, 180))
        tilt_range = self.ptz_config.get('tilt_range', (-45, 45))
        pan_center = self.ptz_config.get('pan_center', 90)
        tilt_center = self.ptz_config.get('tilt_center', 0)
        
        # 将画面比例转换为角度
        # x_ratio=0 对应 pan_range[0], x_ratio=1 对应 pan_range[1]
        pan = pan_range[0] + (pan_range[1] - pan_range[0]) * x_ratio
        
        # y_ratio=0.5 对应 tilt_center, y_ratio=0 对应 tilt_range[0], y_ratio=1 对应 tilt_range[1]
        tilt = tilt_range[0] + (tilt_range[1] - tilt_range[0]) * y_ratio
        
        return (pan, tilt, zoom)
    
    def move_to_target(self, x_ratio: float, y_ratio: float, 
                       zoom: int = None) -> bool:
        """
        移动到目标位置
        Args:
            x_ratio: X方向比例 (0-1)
            y_ratio: Y方向比例 (0-1)
            zoom: 变倍
        Returns:
            是否成功
        """
        pan, tilt, zoom = self.calculate_ptz_position(x_ratio, y_ratio, zoom)
        return self.goto_exact_position(pan, tilt, zoom)
    
    def track_target(self, x_ratio: float, y_ratio: float, 
                     zoom: int = None) -> bool:
        """
        跟踪目标 - 与move_to_target相同，但可以添加跟踪特定逻辑
        Args:
            x_ratio: X方向比例
            y_ratio: Y方向比例
            zoom: 变倍
        Returns:
            是否成功
        """
        return self.move_to_target(x_ratio, y_ratio, zoom)
    
    def get_current_position(self) -> PTZPosition:
        """获取当前位置"""
        with self.position_lock:
            return PTZPosition(
                pan=self.current_position.pan,
                tilt=self.current_position.tilt,
                zoom=self.current_position.zoom
            )


class PTZController:
    """
    PTZ高级控制器
    提供平滑移动、跟踪等功能
    """
    
    def __init__(self, ptz_camera: PTZCamera):
        """
        初始化控制器
        Args:
            ptz_camera: PTZ摄像头实例
        """
        self.ptz = ptz_camera
        self.tracking = False
        self.tracking_thread = None
        self.target_position = None
    
    def smooth_move_to(self, pan: float, tilt: float, zoom: int,
                       steps: int = 10, delay: float = 0.1) -> bool:
        """
        平滑移动到目标位置
        Args:
            pan: 目标水平角度
            tilt: 目标垂直角度
            zoom: 目标变倍
            steps: 移动步数
            delay: 步间延迟
        Returns:
            是否成功
        """
        current = self.ptz.get_current_position()
        
        # 计算步长
        pan_step = (pan - current.pan) / steps
        tilt_step = (tilt - current.tilt) / steps
        zoom_step = (zoom - current.zoom) / steps
        
        for i in range(1, steps + 1):
            current_pan = current.pan + pan_step * i
            current_tilt = current.tilt + tilt_step * i
            current_zoom = int(current.zoom + zoom_step * i)
            
            self.ptz.goto_exact_position(current_pan, current_tilt, current_zoom)
            time.sleep(delay)
        
        return True
    
    def start_tracking(self, get_target_func, update_interval: float = 0.1):
        """
        开始跟踪
        Args:
            get_target_func: 获取目标位置的函数 (返回 x_ratio, y_ratio 或 None)
            update_interval: 更新间隔
        """
        self.tracking = True
        
        def tracking_worker():
            while self.tracking:
                try:
                    target = get_target_func()
                    if target:
                        x_ratio, y_ratio = target
                        self.ptz.track_target(x_ratio, y_ratio)
                    time.sleep(update_interval)
                except Exception as e:
                    print(f"跟踪错误: {e}")
                    time.sleep(0.1)
        
        self.tracking_thread = threading.Thread(target=tracking_worker, daemon=True)
        self.tracking_thread.start()
    
    def stop_tracking(self):
        """停止跟踪"""
        self.tracking = False
        if self.tracking_thread:
            self.tracking_thread.join(timeout=1)
            self.tracking_thread = None
    
    def zoom_to_target_size(self, target_size: Tuple[int, int],
                            frame_size: Tuple[int, int],
                            min_zoom: int = 2, max_zoom: int = 20) -> int:
        """
        根据目标大小计算合适的变倍
        Args:
            target_size: 目标尺寸 (width, height)
            frame_size: 画面尺寸 (width, height)
            min_zoom: 最小变倍
            max_zoom: 最大变倍
        Returns:
            计算的变倍值
        """
        target_area = target_size[0] * target_size[1]
        frame_area = frame_size[0] * frame_size[1]
        
        # 目标占画面比例
        ratio = target_area / frame_area
        
        # 根据比例计算变倍 (目标占画面30%时变倍为1)
        if ratio > 0:
            ideal_zoom = math.sqrt(0.3 / ratio)
            zoom = int(max(min_zoom, min(max_zoom, ideal_zoom)))
        else:
            zoom = min_zoom
        
        return zoom