dsh/dual_camera_system/polling_tracker.py

"""
轮询跟踪 + PTZ 抓拍协调器
"""

import os
import time
import threading
import logging
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

import cv2
import numpy as np

from config import TRACKING_CONFIG
from tracker import UltralyticsTracker, TrackedPerson
from coordinator import TargetSelector

logger = logging.getLogger(__name__)


@dataclass
class CaptureRecord:
    """单次抓拍记录"""
    track_id: int
    timestamp: float
    position: Tuple[float, float]
    ptz_position: Tuple[float, float, int]
    ptz_image: np.ndarray
    panorama_image: Optional[np.ndarray]
    confidence: float


class PollingTrackingCoordinator:
    """多目标轮询跟踪 + PTZ 抓拍协调器"""

    def __init__(
        self,
        panorama_camera,
        ptz_camera,
        tracker: UltralyticsTracker,
        config: Optional[Dict] = None,
        calibrator=None,
    ):
        self.panorama = panorama_camera
        self.ptz = ptz_camera
        self.tracker = tracker
        self.config = config or TRACKING_CONFIG
        self.calibrator = calibrator

        self.active_targets: Dict[int, TrackedPerson] = {}
        self.target_order: List[int] = []
        self.current_index: int = 0
        self.batch_captures: List[CaptureRecord] = []
        self._capture_counts: Dict[int, int] = {}

        self.targets_lock = threading.Lock()
        self.running = False
        self._detection_thread = None
        self._ptz_thread = None
        self._paused = False
        self._paused_event = threading.Event()
        self._paused_event.set()

        self.target_selector = TargetSelector(self.config.get("target_selection", {}))
        self.event_pusher = None

        self.stats = {
            "frames_processed": 0,
            "persons_detected": 0,
            "captures": 0,
            "uploads": 0,
            "start_time": None,
        }
        self.stats_lock = threading.Lock()

        self._ensure_capture_dir()

    def set_event_pusher(self, event_pusher):
        self.event_pusher = event_pusher

    def _ensure_capture_dir(self):
        capture_dir = self.config.get("capture_dir", "/home/admin/dsh/tracking_captures")
        try:
            os.makedirs(capture_dir, exist_ok=True)
        except OSError as e:
            logger.warning(f"无法创建抓拍目录 {capture_dir}: {e}")

    def start(self) -> bool:
        if not self.panorama.connect():
            print("连接全景摄像头失败")
            return False
        if not self.ptz.connect():
            print("连接球机失败")
            self.panorama.disconnect()
            return False
        if not self.panorama.start_stream_rtsp():
            print("启动全景视频流失败")
            self.panorama.disconnect()
            self.ptz.disconnect()
            return False

        self.running = True
        self._detection_thread = threading.Thread(target=self._detection_worker, daemon=True)
        self._detection_thread.start()
        self._ptz_thread = threading.Thread(target=self._ptz_worker, daemon=True)
        self._ptz_thread.start()

        with self.stats_lock:
            self.stats["start_time"] = time.time()

        print("轮询跟踪抓拍协调器已启动")
        return True

    def stop(self):
        self.running = False
        if self._detection_thread:
            self._detection_thread.join(timeout=3)
        if self._ptz_thread:
            self._ptz_thread.join(timeout=3)
        self.panorama.disconnect()
        self.ptz.disconnect()
        print("轮询跟踪抓拍协调器已停止")

    def pause(self):
        self._paused = True
        self._paused_event.clear()

    def resume(self):
        self._paused = False
        self._paused_event.set()

    def _detection_worker(self):
        self._paused_event.wait()
        detection_fps = self.config.get("detection_fps", 2)
        detection_interval = 1.0 / detection_fps
        last_detection_time = 0

        while self.running:
            try:
                frame = self.panorama.get_frame()
                if frame is None:
                    time.sleep(0.01)
                    continue

                self._update_stats("frames_processed")

                current_time = time.time()
                if not self._paused and current_time - last_detection_time >= detection_interval:
                    last_detection_time = current_time
                    tracked = self.tracker.update(frame)
                    self._update_active_targets(tracked, frame.shape)
                    if tracked:
                        self._update_stats("persons_detected", len(tracked))

                time.sleep(0.01)
            except Exception as e:
                logger.error(f"检测线程错误: {e}")
                time.sleep(0.1)

    def _update_active_targets(self, tracked: List[TrackedPerson], frame_shape):
        current_time = time.time()
        frame_h, frame_w = frame_shape[:2]
        timeout = self.config.get("tracking_timeout", 3.0)
        max_targets = self.config.get("max_tracking_targets", 4)

        with self.targets_lock:
            # 更新或新增
            updated_ids = set()
            for p in tracked:
                if p.track_id < 0:
                    continue
                updated_ids.add(p.track_id)
                self.active_targets[p.track_id] = p
                if p.track_id not in self.target_order:
                    self.target_order.append(p.track_id)

            # 标记丢失
            lost_ids = [tid for tid in self.target_order if tid not in updated_ids]
            for tid in lost_ids:
                t = self.active_targets.get(tid)
                if t is not None:
                    t.lost = True

            # 移除长期丢失
            remove_ids = []
            for tid in self.target_order:
                t = self.active_targets.get(tid)
                if t is None:
                    remove_ids.append(tid)
                    continue
                # 简单丢失超时移除
                if t.lost:
                    remove_ids.append(tid)

            for tid in remove_ids:
                if tid in self.target_order:
                    self.target_order.remove(tid)
                self.active_targets.pop(tid, None)
                self._capture_counts.pop(tid, None)

            # 人数上限淘汰
            if len(self.active_targets) > max_targets:
                self._prune_targets(frame_w, frame_h, max_targets)

    def _prune_targets(self, frame_w: int, frame_h: int, max_targets: int):
        targets = list(self.active_targets.values())
        frame_size = (frame_w, frame_h)
        scored = []
        for t in targets:
            target_wrapper = type("T", (), {
                "track_id": t.track_id,
                "area": (t.bbox[2] - t.bbox[0]) * (t.bbox[3] - t.bbox[1]),
                "confidence": t.confidence,
                "center_distance": self._center_distance(t.center, frame_size),
                "score": 0.0,
            })()
            target_wrapper.score = self.target_selector.calculate_score(target_wrapper, frame_size)
            scored.append(target_wrapper)
        scored.sort(key=lambda x: x.score, reverse=True)
        keep_ids = {t.track_id for t in scored[:max_targets]}
        remove_ids = [tid for tid in self.active_targets if tid not in keep_ids]
        for tid in remove_ids:
            self.active_targets.pop(tid, None)
            if tid in self.target_order:
                self.target_order.remove(tid)
            self._capture_counts.pop(tid, None)

    def _center_distance(self, center: Tuple[int, int], frame_size: Tuple[int, int]) -> float:
        cx, cy = frame_size[0] / 2, frame_size[1] / 2
        dx = abs(center[0] - cx) / cx
        dy = abs(center[1] - cy) / cy
        return (dx + dy) / 2

    def _ptz_worker(self):
        while self.running:
            try:
                if self._paused:
                    self._paused_event.wait()
                    continue

                with self.targets_lock:
                    has_targets = bool(self.active_targets)
                    target_order_snapshot = self.target_order.copy()

                if not has_targets or not target_order_snapshot:
                    self._flush_batch_if_needed()
                    time.sleep(0.1)
                    continue

                if self.current_index >= len(target_order_snapshot):
                    self.current_index = 0

                target_id = target_order_snapshot[self.current_index]

                with self.targets_lock:
                    target = self.active_targets.get(target_id)

                if target is None or target.lost:
                    self._advance(len(target_order_snapshot))
                    continue

                record = self._capture_one(target)
                if record:
                    self.batch_captures.append(record)
                    self._update_stats("captures")

                self._advance(len(target_order_snapshot))

                # 一轮完成
                if self.current_index == 0 and self.batch_captures:
                    self._upload_batch(self.batch_captures)
                    self.batch_captures.clear()

            except Exception as e:
                logger.error(f"PTZ 线程错误: {e}")
                time.sleep(0.1)

    def _advance(self, order_len: int = None):
        if order_len is None:
            order_len = len(self.target_order) or 1
        self.current_index = (self.current_index + 1) % order_len

    def _capture_one(self, target: TrackedPerson) -> Optional[CaptureRecord]:
        frame = self.panorama.get_frame()
        if frame is None:
            return None
        frame_h, frame_w = frame.shape[:2]

        x_ratio = target.center[0] / frame_w
        y_ratio = target.center[1] / frame_h

        if self.calibrator and self.calibrator.is_calibrated():
            pan, tilt = self.calibrator.transform(x_ratio, y_ratio)
            ptz_config = getattr(self.ptz, "ptz_config", {})
            if ptz_config.get("pan_flip"):
                pan = (pan + 180) % 360
            zoom = ptz_config.get("default_zoom", 8)
        else:
            pan, tilt, zoom = self.ptz.calculate_ptz_position(x_ratio, y_ratio)

        success = self.ptz.goto_exact_position(pan, tilt, zoom)
        if not success:
            return None

        time.sleep(self.config.get("ptz_stabilize_time", 2.0))
        ptz_frame = self._get_clear_ptz_frame()
        if ptz_frame is None:
            return None

        max_cap = self.config.get("max_capture_per_target", 0)
        if max_cap > 0 and self._capture_counts.get(target.track_id, 0) >= max_cap:
            return None
        self._capture_counts[target.track_id] = self._capture_counts.get(target.track_id, 0) + 1

        panorama_image = frame.copy() if self.config.get("save_panorama_pair", True) else None

        # 本地保存
        self._save_local(ptz_frame, panorama_image, target, pan, tilt, zoom)

        return CaptureRecord(
            track_id=target.track_id,
            timestamp=time.time(),
            position=(x_ratio, y_ratio),
            ptz_position=(pan, tilt, zoom),
            ptz_image=ptz_frame,
            panorama_image=panorama_image,
            confidence=target.confidence,
        )

    def _get_clear_ptz_frame(self, max_attempts: int = 5, wait_interval: float = 0.2) -> Optional[np.ndarray]:
        best_frame = None
        best_score = -1
        for _ in range(max_attempts):
            frame = self.ptz.get_frame()
            if frame is not None:
                frame_copy = frame.copy()
                gray = cv2.cvtColor(frame_copy, cv2.COLOR_BGR2GRAY)
                score = cv2.Laplacian(gray, cv2.CV_64F).var()
                if score > best_score:
                    best_score = score
                    best_frame = frame_copy
            time.sleep(wait_interval)
        return best_frame

    def _save_local(self, ptz_frame, panorama_image, target, pan, tilt, zoom):
        capture_dir = self.config.get("capture_dir", "/home/admin/dsh/tracking_captures")
        try:
            os.makedirs(capture_dir, exist_ok=True)
        except OSError as e:
            logger.warning(f"无法创建抓拍目录 {capture_dir}: {e}")
            return
        timestamp = int(time.time() * 1000)
        base = f"{capture_dir}/ptz_{target.track_id}_{timestamp}_{pan:.0f}_{tilt:.0f}_z{zoom}.jpg"
        cv2.imwrite(base, ptz_frame)
        if panorama_image is not None:
            pan_base = f"{capture_dir}/panorama_{target.track_id}_{timestamp}.jpg"
            cv2.imwrite(pan_base, panorama_image)

    def _upload_batch(self, records: List[CaptureRecord]):
        if not self.event_pusher or not self.config.get("enable_upload", True):
            return
        try:
            uploads = []
            for r in records:
                ptz_url = self.event_pusher.upload_numpy_image(r.ptz_image)
                pan_url = None
                if r.panorama_image is not None:
                    pan_url = self.event_pusher.upload_numpy_image(r.panorama_image)
                uploads.append({
                    "track_id": r.track_id,
                    "ptz_image_url": ptz_url,
                    "panorama_image_url": pan_url,
                    "position": r.position,
                    "ptz_position": r.ptz_position,
                    "confidence": r.confidence,
                    "timestamp": r.timestamp,
                })
            self.event_pusher.push_tracking_capture(batch_time=time.time(), captures=uploads)
            self._update_stats("uploads")
        except Exception as e:
            logger.error(f"批量上传失败: {e}")

    def _flush_batch_if_needed(self):
        if self.batch_captures:
            self._upload_batch(self.batch_captures)
            self.batch_captures.clear()

    def _update_stats(self, key: str, value: int = 1):
        with self.stats_lock:
            if key in self.stats:
                self.stats[key] += value

    def get_stats(self) -> Dict:
        with self.stats_lock:
            return self.stats.copy()