dsh/dual_camera_system/calibration.py

"""
相机校准模块
实现全景相机与球机的自动校准
建立画面坐标到PTZ角度的映射关系
"""

import time
import math
import threading
import numpy as np
import cv2
from typing import List, Tuple, Dict, Optional, Callable
from dataclasses import dataclass, field
from enum import Enum

from ptz_camera import PTZCamera


class CalibrationState(Enum):
    """校准状态"""
    IDLE = 0           # 空闲
    RUNNING = 1        # 校准中
    SUCCESS = 2        # 校准成功
    FAILED = 3         # 校准失败


@dataclass
class CalibrationPoint:
    """校准点"""
    pan: float              # 球机水平角度
    tilt: float             # 球机垂直角度
    zoom: float             # 变倍
    x_ratio: float = 0.0    # 全景画面X比例
    y_ratio: float = 0.0    # 全景画面Y比例
    detected: bool = False  # 是否已检测到


@dataclass
class CalibrationResult:
    """校准结果"""
    success: bool
    points: List[CalibrationPoint]
    transform_matrix: Optional[np.ndarray] = None  # 变换矩阵
    error_message: str = ""
    rms_error: float = 0.0  # 均方根误差


class VisualCalibrationDetector:
    """
    视觉校准检测器
    通过视觉方法检测球机在全景画面中的位置
    """
    
    def __init__(self):
        """初始化检测器"""
        # 特征匹配器 - 优先SIFT，备选ORB
        try:
            self.feature_detector = cv2.SIFT_create()
            self.feature_type = 'SIFT'
        except AttributeError:
            # SIFT需要opencv-contrib，ORB是基础版自带的
            self.feature_detector = cv2.ORB_create(nfeatures=500)
            self.feature_type = 'ORB'
            print("提示: 使用ORB特征检测器 (安装opencv-contrib可启用SIFT)")
        
        self.matcher = cv2.BFMatcher(cv2.NORM_L2 if self.feature_type == 'SIFT' else cv2.NORM_HAMMING)
        
        # 校准模式
        self.use_motion_detection = True
        self.use_feature_matching = True
    
    def detect_by_motion(self, frames_before: np.ndarray, 
                         frames_after: np.ndarray) -> Optional[Tuple[float, float]]:
        """
        通过运动检测定位球机指向位置
        原理: 球机移动会在全景画面中产生运动区域
        
        Args:
            frames_before: 球机移动前的全景帧 (多帧平均)
            frames_after: 球机移动后的全景帧 (多帧平均)
        Returns:
            (x_ratio, y_ratio) 或 None
        """
        if frames_before is None or frames_after is None:
            return None
        
        # 计算帧差
        if len(frames_before.shape) == 3:
            before_gray = cv2.cvtColor(frames_before, cv2.COLOR_BGR2GRAY)
        else:
            before_gray = frames_before
            
        if len(frames_after.shape) == 3:
            after_gray = cv2.cvtColor(frames_after, cv2.COLOR_BGR2GRAY)
        else:
            after_gray = frames_after
        
        # 计算差异
        diff = cv2.absdiff(before_gray, after_gray)
        
        # 二值化
        _, thresh = cv2.threshold(diff, 25, 255, cv2.THRESH_BINARY)
        
        # 形态学操作去噪
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
        
        # 查找轮廓
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if not contours:
            return None
        
        # 找到最大的运动区域
        max_contour = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(max_contour)
        
        # 过滤太小的区域
        if area < 500:  # 最小面积阈值
            return None
        
        # 计算中心点
        M = cv2.moments(max_contour)
        if M["m00"] == 0:
            return None
        
        cx = M["m10"] / M["m00"]
        cy = M["m01"] / M["m00"]
        
        # 转换为比例
        h, w = before_gray.shape
        x_ratio = cx / w
        y_ratio = cy / h
        
        print(f"      运动检测: 中心=({cx:.1f}, {cy:.1f}), 面积={area:.0f}")
        
        return (x_ratio, y_ratio)
    
    def detect_by_feature_match(self, panorama_frame: np.ndarray,
                                 ptz_frame: np.ndarray) -> Optional[Tuple[float, float]]:
        """
        通过特征匹配定位
        将球机抓拍的图像与全景画面进行特征匹配
        
        Args:
            panorama_frame: 全景画面
            ptz_frame: 球机抓拍画面
        Returns:
            (x_ratio, y_ratio) 或 None
        """
        if panorama_frame is None or ptz_frame is None:
            return None
        
        # 转换为灰度
        if len(panorama_frame.shape) == 3:
            pan_gray = cv2.cvtColor(panorama_frame, cv2.COLOR_BGR2GRAY)
        else:
            pan_gray = panorama_frame
            
        if len(ptz_frame.shape) == 3:
            ptz_gray = cv2.cvtColor(ptz_frame, cv2.COLOR_BGR2GRAY)
        else:
            ptz_gray = ptz_frame
        
        # 检测特征点
        try:
            kp1, des1 = self.feature_detector.detectAndCompute(ptz_gray, None)
            kp2, des2 = self.feature_detector.detectAndCompute(pan_gray, None)
            
            if des1 is None or des2 is None:
                return None
            
            if len(kp1) < 4 or len(kp2) < 4:
                return None
            
            # 特征匹配
            matches = self.matcher.knnMatch(des1, des2, k=2)
            
            # 筛选好的匹配
            good_matches = []
            for m, n in matches:
                if m.distance < 0.75 * n.distance:
                    good_matches.append(m)
            
            if len(good_matches) < 4:
                return None
            
            # 获取匹配点坐标
            ptz_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches])
            pan_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches])
            
            # 计算全景画面中匹配点的中心
            center_x = np.mean(pan_pts[:, 0])
            center_y = np.mean(pan_pts[:, 1])
            
            # 转换为比例
            h, w = pan_gray.shape
            x_ratio = center_x / w
            y_ratio = center_y / h
            
            print(f"      特征匹配: 匹配点={len(good_matches)}, 中心=({center_x:.1f}, {center_y:.1f})")
            
            return (x_ratio, y_ratio)
            
        except Exception as e:
            print(f"      特征匹配错误: {e}")
            return None
    
    def detect_position(self, panorama_frame: np.ndarray,
                        frames_before: np.ndarray = None,
                        frames_after: np.ndarray = None,
                        ptz_frame: np.ndarray = None) -> Tuple[bool, float, float]:
        """
        综合检测球机在全景画面中的位置
        
        Args:
            panorama_frame: 全景画面
            frames_before: 移动前的帧 (用于运动检测)
            frames_after: 移动后的帧 (用于运动检测)
            ptz_frame: 球机抓拍画面 (用于特征匹配)
        Returns:
            (是否成功, x_ratio, y_ratio)
        """
        results = []
        
        # 方法1: 运动检测
        if self.use_motion_detection and frames_before is not None and frames_after is not None:
            motion_result = self.detect_by_motion(frames_before, frames_after)
            if motion_result:
                results.append(('motion', motion_result, 0.6))  # 权重0.6
        
        # 方法2: 特征匹配
        if self.use_feature_matching and ptz_frame is not None:
            feature_result = self.detect_by_feature_match(panorama_frame, ptz_frame)
            if feature_result:
                results.append(('feature', feature_result, 0.4))  # 权重0.4
        
        if not results:
            return (False, 0.0, 0.0)
        
        # 加权融合结果
        total_weight = sum(r[2] for r in results)
        x_sum = sum(r[1][0] * r[2] for r in results)
        y_sum = sum(r[1][1] * r[2] for r in results)
        
        x_ratio = x_sum / total_weight
        y_ratio = y_sum / total_weight
        
        print(f"      融合结果: ({x_ratio:.3f}, {y_ratio:.3f})")
        
        return (True, x_ratio, y_ratio)


class CameraCalibrator:
    """
    相机校准器
    建立全景相机坐标与球机PTZ角度的映射关系
    使用视觉方法进行精确校准
    """
    
    def __init__(self, ptz_camera: PTZCamera, 
                 get_frame_func: Callable[[], np.ndarray],
                 detect_marker_func: Callable[[np.ndarray], Optional[Tuple[float, float]]] = None,
                 ptz_capture_func: Callable[[], np.ndarray] = None):
        """
        初始化校准器
        Args:
            ptz_camera: 球机实例
            get_frame_func: 获取全景画面的函数
            detect_marker_func: 检测标记点的函数 (返回x_ratio, y_ratio)
            ptz_capture_func: 球机抓拍函数 (用于特征匹配)
        """
        self.ptz = ptz_camera
        self.get_frame = get_frame_func
        self.detect_marker = detect_marker_func
        self.ptz_capture = ptz_capture_func
        
        # 视觉检测器
        self.visual_detector = VisualCalibrationDetector()
        
        self.state = CalibrationState.IDLE
        self.result: Optional[CalibrationResult] = None
        
        # 校准点配置 (覆盖画面不同区域)
        self.calibration_points = self._generate_calibration_points()
        
        # 变换参数
        self.pan_offset = 0.0      # 水平偏移
        self.tilt_offset = 0.0     # 垂直偏移
        self.pan_scale = 1.0       # 水平缩放
        self.tilt_scale = 1.0      # 垂直缩放
        
        # 校准配置
        self.stabilize_time = 2.0   # 球机稳定等待时间
        self.use_motion_detection = True
        self.use_feature_matching = True
        
        # 回调
        self.on_progress: Optional[Callable[[int, int, str], None]] = None
        self.on_complete: Optional[Callable[[CalibrationResult], None]] = None
    
    def _generate_calibration_points(self) -> List[CalibrationPoint]:
        """
        生成校准点
        在球机视野范围内均匀分布
        Returns:
            校准点列表
        """
        points = []
        
        # 5x5网格校准点 (共25个点)
        # 水平角度范围: 0-180度 (全景相机通常覆盖180度)
        # 垂直角度范围: -30到30度
        pan_range = (0, 180)
        tilt_range = (-30, 30)
        
        # 使用较少的点进行快速校准
        grid_size = 3  # 3x3 = 9个点
        
        for i in range(grid_size):
            for j in range(grid_size):
                pan = pan_range[0] + (pan_range[1] - pan_range[0]) * i / (grid_size - 1)
                tilt = tilt_range[0] + (tilt_range[1] - tilt_range[0]) * j / (grid_size - 1)
                
                points.append(CalibrationPoint(
                    pan=pan,
                    tilt=tilt,
                    zoom=1.0  # 校准时使用最小变倍
                ))
        
        return points
    
    def calibrate(self, quick_mode: bool = True) -> CalibrationResult:
        """
        执行校准 - 使用视觉检测方法
        Args:
            quick_mode: 快速模式 (使用较少校准点)
        Returns:
            校准结果
        """
        self.state = CalibrationState.RUNNING
        
        if quick_mode:
            calib_points = self._get_quick_calibration_points()
        else:
            calib_points = self.calibration_points
        
        total_points = len(calib_points)
        valid_points = []
        
        print(f"开始视觉校准, 共 {total_points} 个校准点...")
        print(f"检测方法: 运动检测={'开启' if self.use_motion_detection else '关闭'}, "
              f"特征匹配={'开启' if self.use_feature_matching else '关闭'}")
        
        for idx, point in enumerate(calib_points):
            if self.on_progress:
                self.on_progress(idx + 1, total_points, 
                               f"正在校准点 {idx + 1}/{total_points}: pan={point.pan:.1f}°, tilt={point.tilt:.1f}°")
            
            print(f"  校准点 {idx + 1}/{total_points}: pan={point.pan:.1f}°, tilt={point.tilt:.1f}°")
            
            # ===== 步骤1: 获取移动前的全景帧 =====
            print(f"    [1/4] 获取移动前的全景帧...")
            frames_before_list = []
            for _ in range(3):  # 取3帧平均
                frame = self.get_frame()
                if frame is not None:
                    frames_before_list.append(frame)
                time.sleep(0.1)
            
            if not frames_before_list:
                print(f"    警告: 无法获取移动前的全景画面")
                continue
            
            frames_before = np.mean(frames_before_list, axis=0).astype(np.uint8)
            
            # ===== 步骤2: 移动球机到指定位置 =====
            print(f"    [2/4] 移动球机到目标位置...")
            if not self.ptz.goto_exact_position(point.pan, point.tilt, 1):
                print(f"    警告: 移动球机失败")
                continue
            
            # 等待球机稳定
            time.sleep(self.stabilize_time)
            
            # ===== 步骤3: 获取移动后的全景帧和球机抓拍 =====
            print(f"    [3/4] 获取移动后的帧...")
            frames_after_list = []
            for _ in range(3):
                frame = self.get_frame()
                if frame is not None:
                    frames_after_list.append(frame)
                time.sleep(0.1)
            
            if not frames_after_list:
                print(f"    警告: 无法获取移动后的全景画面")
                continue
            
            frames_after = np.mean(frames_after_list, axis=0).astype(np.uint8)
            panorama_frame = frames_after  # 当前全景帧
            
            # 球机抓拍 (用于特征匹配)
            ptz_frame = None
            if self.use_feature_matching and self.ptz_capture:
                try:
                    ptz_frame = self.ptz_capture()
                    if ptz_frame is not None:
                        print(f"      球机抓拍成功: {ptz_frame.shape}")
                except Exception as e:
                    print(f"      球机抓拍失败: {e}")
            
            # ===== 步骤4: 视觉检测位置 =====
            print(f"    [4/4] 视觉检测位置...")
            
            # 优先使用自定义标记检测
            if self.detect_marker:
                marker_pos = self.detect_marker(panorama_frame)
                if marker_pos:
                    point.x_ratio, point.y_ratio = marker_pos
                    point.detected = True
                    valid_points.append(point)
                    print(f"    标记检测成功: ({point.x_ratio:.3f}, {point.y_ratio:.3f})")
                    continue
            
            # 使用视觉检测器
            detected, x_ratio, y_ratio = self.visual_detector.detect_position(
                panorama_frame=panorama_frame,
                frames_before=frames_before,
                frames_after=frames_after,
                ptz_frame=ptz_frame
            )
            
            if detected:
                point.x_ratio = x_ratio
                point.y_ratio = y_ratio
                point.detected = True
                valid_points.append(point)
                print(f"    ✓ 视觉检测成功: ({point.x_ratio:.3f}, {point.y_ratio:.3f})")
            else:
                # 降级: 使用估算方法
                print(f"    视觉检测失败, 使用估算方法...")
                x_ratio = point.pan / 180.0
                y_ratio = 0.5 + point.tilt / 90.0
                x_ratio = max(0, min(1, x_ratio))
                y_ratio = max(0, min(1, y_ratio))
                point.x_ratio = x_ratio
                point.y_ratio = y_ratio
                point.detected = True
                valid_points.append(point)
                print(f"    △ 使用估算: ({point.x_ratio:.3f}, {point.y_ratio:.3f})")
        
        # 检查校准点数量
        if len(valid_points) < 4:
            self.state = CalibrationState.FAILED
            self.result = CalibrationResult(
                success=False,
                points=valid_points,
                error_message=f"有效校准点不足 (需要至少4个, 实际{len(valid_points)}个)"
            )
            print(f"校准失败: {self.result.error_message}")
            
            if self.on_complete:
                self.on_complete(self.result)
            
            return self.result
        
        # 计算变换参数
        success = self._calculate_transform(valid_points)
        
        if success:
            self.state = CalibrationState.SUCCESS
            rms_error = self._calculate_rms_error(valid_points)
            
            self.result = CalibrationResult(
                success=True,
                points=valid_points,
                rms_error=rms_error
            )
            print(f"校准成功! RMS误差: {rms_error:.4f}")
        else:
            self.state = CalibrationState.FAILED
            self.result = CalibrationResult(
                success=False,
                points=valid_points,
                error_message="变换参数计算失败"
            )
            print(f"校准失败: {self.result.error_message}")
        
        if self.on_complete:
            self.on_complete(self.result)
        
        return self.result
    
    def _get_quick_calibration_points(self) -> List[CalibrationPoint]:
        """获取快速校准点 (5个点)"""
        return [
            CalibrationPoint(pan=0, tilt=0, zoom=1),        # 左侧
            CalibrationPoint(pan=90, tilt=0, zoom=1),      # 中心
            CalibrationPoint(pan=180, tilt=0, zoom=1),     # 右侧
            CalibrationPoint(pan=90, tilt=-20, zoom=1),    # 上方
            CalibrationPoint(pan=90, tilt=20, zoom=1),     # 下方
        ]
    
    def _calculate_transform(self, points: List[CalibrationPoint]) -> bool:
        """
        计算坐标变换参数
        Args:
            points: 有效校准点
        Returns:
            是否成功
        """
        try:
            # 使用最小二乘法拟合变换关系
            # 简单线性模型: pan = a0 + a1*x + a2*y
            #             tilt = b0 + b1*x + b2*y
            
            n = len(points)
            
            # 构建矩阵
            A = np.ones((n, 3))
            A[:, 1] = [p.x_ratio for p in points]
            A[:, 2] = [p.y_ratio for p in points]
            
            pan_values = np.array([p.pan for p in points])
            tilt_values = np.array([p.tilt for p in points])
            
            # 最小二乘求解
            pan_params, _, _, _ = np.linalg.lstsq(A, pan_values, rcond=None)
            tilt_params, _, _, _ = np.linalg.lstsq(A, tilt_values, rcond=None)
            
            # 存储变换参数
            self.pan_offset = pan_params[0]
            self.pan_scale_x = pan_params[1]
            self.pan_scale_y = pan_params[2]
            
            self.tilt_offset = tilt_params[0]
            self.tilt_scale_x = tilt_params[1]
            self.tilt_scale_y = tilt_params[2]
            
            print(f"变换参数:")
            print(f"  pan = {self.pan_offset:.2f} + {self.pan_scale_x:.2f}*x + {self.pan_scale_y:.2f}*y")
            print(f"  tilt = {self.tilt_offset:.2f} + {self.tilt_scale_x:.2f}*x + {self.tilt_scale_y:.2f}*y")
            
            return True
            
        except Exception as e:
            print(f"计算变换参数错误: {e}")
            return False
    
    def _calculate_rms_error(self, points: List[CalibrationPoint]) -> float:
        """计算均方根误差"""
        total_error = 0.0
        
        for p in points:
            pred_pan, pred_tilt = self.transform(p.x_ratio, p.y_ratio)
            error = math.sqrt((pred_pan - p.pan)**2 + (pred_tilt - p.tilt)**2)
            total_error += error**2
        
        return math.sqrt(total_error / len(points))
    
    def transform(self, x_ratio: float, y_ratio: float) -> Tuple[float, float]:
        """
        将全景坐标转换为PTZ角度
        Args:
            x_ratio: X方向比例 (0-1)
            y_ratio: Y方向比例 (0-1)
        Returns:
            (pan, tilt) 角度
        """
        pan = self.pan_offset + self.pan_scale_x * x_ratio + self.pan_scale_y * y_ratio
        tilt = self.tilt_offset + self.tilt_scale_x * x_ratio + self.tilt_scale_y * y_ratio
        return (pan, tilt)
    
    def inverse_transform(self, pan: float, tilt: float) -> Tuple[float, float]:
        """
        将PTZ角度转换为全景坐标 (近似)
        Args:
            pan: 水平角度
            tilt: 垂直角度
        Returns:
            (x_ratio, y_ratio)
        """
        # 简化逆变换
        x_ratio = (pan - self.pan_offset) / self.pan_scale_x if self.pan_scale_x != 0 else 0.5
        y_ratio = (tilt - self.tilt_offset) / self.tilt_scale_y if self.tilt_scale_y != 0 else 0.5
        
        return (max(0, min(1, x_ratio)), max(0, min(1, y_ratio)))
    
    def is_calibrated(self) -> bool:
        """是否已完成校准"""
        return self.state == CalibrationState.SUCCESS
    
    def get_state(self) -> CalibrationState:
        """获取当前状态"""
        return self.state
    
    def get_result(self) -> Optional[CalibrationResult]:
        """获取校准结果"""
        return self.result
    
    def save_calibration(self, filepath: str) -> bool:
        """
        保存校准结果
        Args:
            filepath: 文件路径
        Returns:
            是否成功
        """
        if not self.is_calibrated():
            return False
        
        try:
            import json
            data = {
                'pan_offset': self.pan_offset,
                'pan_scale_x': self.pan_scale_x,
                'pan_scale_y': self.pan_scale_y,
                'tilt_offset': self.tilt_offset,
                'tilt_scale_x': self.tilt_scale_x,
                'tilt_scale_y': self.tilt_scale_y,
                'rms_error': self.result.rms_error if self.result else 0,
            }
            
            with open(filepath, 'w') as f:
                json.dump(data, f, indent=2)
            
            print(f"校准结果已保存: {filepath}")
            return True
            
        except Exception as e:
            print(f"保存校准结果失败: {e}")
            return False
    
    def load_calibration(self, filepath: str) -> bool:
        """
        加载校准结果
        Args:
            filepath: 文件路径
        Returns:
            是否成功
        """
        try:
            import json
            
            with open(filepath, 'r') as f:
                data = json.load(f)
            
            self.pan_offset = data['pan_offset']
            self.pan_scale_x = data['pan_scale_x']
            self.pan_scale_y = data['pan_scale_y']
            self.tilt_offset = data['tilt_offset']
            self.tilt_scale_x = data['tilt_scale_x']
            self.tilt_scale_y = data['tilt_scale_y']
            
            self.state = CalibrationState.SUCCESS
            self.result = CalibrationResult(
                success=True,
                points=[],
                rms_error=data.get('rms_error', 0)
            )
            
            print(f"校准结果已加载: {filepath}")
            return True
            
        except FileNotFoundError:
            print(f"校准文件不存在: {filepath}")
            return False
        except Exception as e:
            print(f"加载校准结果失败: {e}")
            return False


class CalibrationManager:
    """
    校准管理器
    管理校准流程和状态
    """
    
    def __init__(self, calibrator: CameraCalibrator):
        """
        初始化管理器
        Args:
            calibrator: 校准器实例
        """
        self.calibrator = calibrator
        self.calibration_file = "calibration.json"
    
    def auto_calibrate(self, force: bool = False) -> CalibrationResult:
        """
        自动校准
        Args:
            force: 是否强制重新校准
        Returns:
            校准结果
        """
        # 尝试加载已有校准结果
        if not force:
            if self.calibrator.load_calibration(self.calibration_file):
                print("使用已有校准结果")
                return self.calibrator.get_result()
        
        # 执行校准
        print("开始自动校准...")
        result = self.calibrator.calibrate(quick_mode=True)
        
        if result.success:
            # 保存校准结果
            self.calibrator.save_calibration(self.calibration_file)
        
        return result
    
    def check_calibration(self) -> Tuple[bool, str]:
        """
        检查校准状态
        Returns:
            (是否有效, 状态信息)
        """
        state = self.calibrator.get_state()
        
        if state == CalibrationState.SUCCESS:
            result = self.calibrator.get_result()
            return (True, f"校准有效, RMS误差: {result.rms_error:.4f}")
        elif state == CalibrationState.FAILED:
            return (False, "校准失败")
        elif state == CalibrationState.RUNNING:
            return (False, "校准进行中")
        else:
            return (False, "未校准")