dsh/dual_camera_system/core/calibration.py

"""加载并应用 PTZ 校准映射."""
import json
import logging
import math
from bisect import bisect_left
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


class CalibrationMapper:
    """
    把视觉坐标 (pan, tilt) 映射到球机物理坐标 (device_pan, device_tilt)。

    支持两种校准文件格式：
    1. 线性/仿射系数格式：
       pan = pan_offset + pan_scale_x * x_ratio + pan_scale_y * y_ratio
       tilt = tilt_offset + tilt_scale_x * x_ratio + tilt_scale_y * y_ratio
    2. 查找表格式：
       pan_lookup: [[x_ratio, pan_angle], ...]
       tilt_lookup: [[y_ratio, tilt_angle], ...]

    x_ratio / y_ratio 由视觉角度和视野配置计算：
       x_ratio = (visual_pan - pan_center) / (pan_range[1] - pan_range[0]) + 0.5
       y_ratio = (visual_tilt - tilt_center) / (tilt_range[1] - tilt_range[0]) + 0.5
    """

    def __init__(
        self,
        calibration_path: Optional[str],
        ptz_config: Dict,
    ):
        self.path = calibration_path
        self.data: Optional[Dict] = None
        self.mount_type: Optional[str] = None
        self.pan_flip: bool = False
        self.tilt_flip: bool = False

        # 视野配置
        self.pan_range = tuple(ptz_config.get("pan_range", (-90, 90)))
        self.pan_center = float(ptz_config.get("pan_center", 0.0))
        self.tilt_range = tuple(ptz_config.get("tilt_range", (-5, 20)))
        self.tilt_center = float(ptz_config.get("tilt_center", 0.0))

        if calibration_path:
            self._load(calibration_path)

    def is_loaded(self) -> bool:
        return self.data is not None

    def _load(self, path: str) -> None:
        try:
            p = Path(path)
            if not p.exists():
                logger.warning("[calibration] 校准文件不存在: %s", path)
                return
            self.data = json.loads(p.read_text(encoding="utf-8"))
            self.mount_type = self.data.get("mount_type")
            self.pan_flip = bool(self.data.get("pan_flip", False))
            self.tilt_flip = bool(self.data.get("tilt_flip", False))
            logger.info("[calibration] 加载校准文件: %s", path)
        except Exception as exc:
            logger.error("[calibration] 加载校准文件失败: %s, %s", path, exc)
            self.data = None

    def _visual_to_ratio(self, pan: float, tilt: float) -> Tuple[float, float]:
        """视觉角度 -> 归一化坐标 (x_ratio, y_ratio)。"""
        pan_span = self.pan_range[1] - self.pan_range[0]
        tilt_span = self.tilt_range[1] - self.tilt_range[0]
        x = 0.5 if pan_span == 0 else (pan - self.pan_center) / pan_span + 0.5
        y = 0.5 if tilt_span == 0 else (tilt - self.tilt_center) / tilt_span + 0.5
        return float(x), float(y)

    @staticmethod
    def _interpolate(lookup: List[List[float]], ratio: float) -> float:
        """分段线性插值。"""
        if not lookup:
            return float("nan")
        ratios = [row[0] for row in lookup]
        values = [row[1] for row in lookup]
        if ratio <= ratios[0]:
            return values[0]
        if ratio >= ratios[-1]:
            return values[-1]
        idx = bisect_left(ratios, ratio)
        x0, x1 = ratios[idx - 1], ratios[idx]
        y0, y1 = values[idx - 1], values[idx]
        if abs(x1 - x0) < 1e-9:
            return y0
        return y0 + (y1 - y0) * (ratio - x0) / (x1 - x0)

    def visual_to_device(self, pan: float, tilt: float) -> Tuple[float, float]:
        """
        把视觉坐标转换为球机物理坐标。

        Returns:
            (device_pan, device_tilt)
        """
        if not self.data:
            return float(pan), float(tilt)

        x, y = self._visual_to_ratio(pan, tilt)

        # 优先使用仿射系数
        pan_offset = self.data.get("pan_offset")
        pan_scale_x = self.data.get("pan_scale_x")
        tilt_offset = self.data.get("tilt_offset")
        tilt_scale_y = self.data.get("tilt_scale_y")

        if (
            pan_offset is not None
            and pan_scale_x is not None
            and tilt_offset is not None
            and tilt_scale_y is not None
        ):
            pan_scale_y = self.data.get("pan_scale_y", 0.0) or 0.0
            tilt_scale_x = self.data.get("tilt_scale_x", 0.0) or 0.0
            device_pan = pan_offset + pan_scale_x * x + pan_scale_y * y
            device_tilt = tilt_offset + tilt_scale_x * x + tilt_scale_y * y
            return float(device_pan), float(device_tilt)

        # 回退到查找表
        pan_lookup = self.data.get("pan_lookup", [])
        tilt_lookup = self.data.get("tilt_lookup", [])
        device_pan = self._interpolate(pan_lookup, x)
        device_tilt = self._interpolate(tilt_lookup, y)
        if math.isnan(device_pan):
            device_pan = pan
        if math.isnan(device_tilt):
            device_tilt = tilt
        return float(device_pan), float(device_tilt)