dsh/dual_camera_system/ptz_camera.py

"""
球机(PTZ)控制模块
负责PTZ控制、精确定位和视频流获取
"""

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

import cv2
import numpy as np

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()
        
        # 视频流 (用于校准时抓拍球机画面)
        self.rtsp_cap = None
        self.current_frame = None
        self.frame_lock = threading.Lock()
        self.stream_thread = None
        self.running_stream = False
    
    def connect(self) -> bool:
        """
        连接球机
        Returns:
            是否成功
        """
        print(f"[PTZCamera] 正在连接球机: IP={self.config['ip']}:{self.config['port']}, 通道={self.config['channel']}")
        
        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"[PTZCamera] 连接球机失败: IP={self.config['ip']}, 错误码={error}")
            print(f"[PTZCamera] 请检查: 1)IP地址是否正确 2)网络是否连通 3)用户名密码是否正确")
            return False
        
        self.login_handle = login_handle
        self.connected = True
        print(f"[PTZCamera] 成功连接球机: {self.config['ip']}, handle={login_handle}")
        print(f"[PTZCamera] 配置: 通道={self.config['channel']}, 默认变倍={self.ptz_config['default_zoom']}")
        return True
    
    def disconnect(self):
        """断开连接"""
        self.stop_stream()
        if self.login_handle:
            self.sdk.logout(self.login_handle)
            self.login_handle = None
        self.connected = False
    
    def start_stream_rtsp(self, rtsp_url: str = None) -> bool:
        """启动RTSP视频流 (用于校准时获取球机画面)"""
        if rtsp_url is None:
            rtsp_url = self.config.get('rtsp_url') or \
                f"rtsp://{self.config['username']}:{self.config['password']}@{self.config['ip']}:{self.config.get('rtsp_port', 554)}/cam/realmonitor?channel=1&subtype=1"
        
        try:
            os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = 'threads;1'
            self.rtsp_cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG)
            if not self.rtsp_cap.isOpened():
                print(f"[PTZCamera] 无法打开RTSP流: {rtsp_url}")
                return False
            
            self.rtsp_cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
            self.running_stream = True
            self.stream_thread = threading.Thread(target=self._stream_worker, daemon=True)
            self.stream_thread.start()
            print(f"[PTZCamera] RTSP视频流已启动")
            return True
        except Exception as e:
            print(f"[PTZCamera] RTSP流启动失败: {e}")
            return False
    
    def _stream_worker(self):
        """视频流工作线程"""
        while self.running_stream:
            try:
                if self.rtsp_cap is None or not self.rtsp_cap.isOpened():
                    time.sleep(0.1)
                    continue
                
                ret, frame = self.rtsp_cap.read()
                if not ret or frame is None:
                    time.sleep(0.01)
                    continue
                
                with self.frame_lock:
                    self.current_frame = frame.copy()
                
                time.sleep(0.001)
            except Exception as e:
                print(f"[PTZCamera] 视频流错误: {e}")
                time.sleep(0.1)
    
    def get_frame(self) -> Optional[np.ndarray]:
        """获取球机当前帧"""
        with self.frame_lock:
            return self.current_frame.copy() if self.current_frame is not None else None
    
    def stop_stream(self):
        """停止视频流"""
        self.running_stream = False
        if self.stream_thread:
            self.stream_thread.join(timeout=2)
            self.stream_thread = None
        if self.rtsp_cap:
            self.rtsp_cap.release()
            self.rtsp_cap = None
    
    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:
            print(f"[PTZCamera] PTZ控制失败: 未连接球机")
            return False
        
        if self.login_handle is None or self.login_handle <= 0:
            print(f"[PTZCamera] PTZ控制失败: 登录句柄无效 (handle={self.login_handle})")
            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:
            是否成功
        """
        param1 = int(pan * 10)
        param2 = int(tilt * 10)
        param3 = int(min(zoom, 128))
        
        print(f"[PTZCamera] goto_exact_position: pan={pan:.1f}° tilt={tilt:.1f}° zoom={zoom} → p1={param1} p2={param2} p3={param3}")
        
        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)
        else:
            print(f"[PTZCamera] goto_exact_position FAILED!")
        
        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