dsh/dual_camera_system/core/spatial_scanner.py

"""球机 360° 扫描与全景图生成."""
import logging
import os
import time
from typing import Callable, Dict, List, Optional, Tuple
import cv2
import numpy as np

logger = logging.getLogger(__name__)

from core.coord_utils import compute_sample_grid


class SpatialScanner:
    def __init__(
        self,
        group_id: str,
        ptz_camera,
        ptz_frame_source: Callable[[], Optional[np.ndarray]],
        data_dir: str = "data",
        stabilize_time: float = 1.5,
    ):
        self.group_id = group_id
        self.ptz = ptz_camera
        self.get_ptz_frame = ptz_frame_source
        self.data_dir = os.path.join(data_dir, group_id)
        os.makedirs(os.path.join(self.data_dir, "samples"), exist_ok=True)
        os.makedirs(os.path.join(self.data_dir, "panorama"), exist_ok=True)
        self.stabilize_time = stabilize_time
        self.progress = {"total": 0, "current": 0, "state": "idle"}
        self.cancelled = False

    def cancel(self):
        self.cancelled = True

    def run(
        self,
        pan_range: Tuple[float, float] = (0.0, 360.0),
        tilt_layers: Tuple[float, ...] = (-20.0, 0.0, 20.0),
        pan_step: float = 30.0,
        zoom: int = 1,
        progress_callback: Optional[Callable[[Dict], None]] = None,
    ) -> Dict:
        config = {
            "pan_range": pan_range,
            "tilt_layers": list(tilt_layers),
            "pan_step": pan_step,
            "zoom": zoom,
        }

        if self.cancelled:
            self.progress = {"total": 0, "current": 0, "state": "cancelled"}
            return {
                "group_id": self.group_id,
                "samples": [],
                "panorama_path": None,
                "config": config,
            }

        grid = compute_sample_grid(pan_range, tilt_layers, pan_step)
        self.progress = {"total": len(grid), "current": 0, "state": "scanning"}
        samples: List[Dict] = []

        for idx, (pan, tilt) in enumerate(grid):
            if self.cancelled:
                break
            try:
                self.ptz.goto_exact_position(pan, tilt, zoom)
                time.sleep(self.stabilize_time)
                frame = self._wait_frame(timeout=5.0)
                if frame is None:
                    continue
                filename = f"p{pan:.1f}_t{tilt:.1f}.jpg"
                path = os.path.join(self.data_dir, "samples", filename)
                cv2.imwrite(path, frame)
                samples.append({
                    "id": idx + 1,
                    "pan": pan,
                    "tilt": tilt,
                    "zoom": zoom,
                    "thumbnail": path,
                })
                self.progress["current"] = len(samples)
                if progress_callback:
                    progress_callback(dict(self.progress))
            except Exception as exc:
                logger.error("Error processing scan sample pan=%s tilt=%s: %s", pan, tilt, exc)
                continue

        panorama_path = self._build_equirectangular(samples, pan_range, tilt_layers)
        self.progress["state"] = "cancelled" if self.cancelled else "done"
        return {
            "group_id": self.group_id,
            "samples": samples,
            "panorama_path": panorama_path,
            "config": config,
        }

    def _wait_frame(self, timeout: float = 5.0) -> Optional[np.ndarray]:
        deadline = time.time() + timeout
        while time.time() < deadline:
            frame = self.get_ptz_frame()
            if frame is not None:
                return frame
            time.sleep(0.1)
        return None

    def _build_equirectangular(
        self,
        samples: List[Dict],
        pan_range: Tuple[float, float],
        tilt_layers: Tuple[float, ...],
        width: int = 4096,
        height: int = 2048,
    ) -> Optional[str]:
        if not samples:
            return None
        panorama = np.zeros((height, width, 3), dtype=np.uint8)
        pan_start, pan_end = pan_range
        for s in samples:
            frame = cv2.imread(s["thumbnail"])
            if frame is None:
                continue
            h, w = frame.shape[:2]
            pan = s["pan"]
            tilt = s["tilt"]
            u_center = int(((pan - pan_start) / (pan_end - pan_start)) * width)
            v_center = int(((90 - tilt) / 180) * height)
            y_start = max(0, v_center - h // 2)
            y_end = min(height, v_center + h // 2)
            x_start = max(0, u_center - w // 2)
            x_end = min(width, u_center + w // 2)
            paste_h = y_end - y_start
            paste_w = x_end - x_start
            if paste_h <= 0 or paste_w <= 0:
                continue
            roi = frame[(h - paste_h) // 2:(h - paste_h) // 2 + paste_h,
                        (w - paste_w) // 2:(w - paste_w) // 2 + paste_w]
            panorama[y_start:y_end, x_start:x_end] = roi

        path = os.path.join(self.data_dir, "panorama", f"scan_{int(time.time())}.jpg")
        cv2.imwrite(path, panorama)
        return path