Chameleon/ptz_tracker_modular/cpp_bridge.py

from __future__ import annotations
import asyncio
import base64
import logging
import time
from typing import List, Tuple

import numpy as np
import cv2
import websockets

from config import CPP_WS_URI, CPP_CTRL_WS_URI, DROP_DECODE_WHEN_BUSY, PUBLISH_RAW_BEFORE_INFER, PREVIEW_TARGET_W
from utils import json_loads, json_dumps
from state import detected_cameras, video_clients, ptz_states, detection_queue
from media import publish_preview, maybe_downscale
from detection import batcher

logger = logging.getLogger("PTZTracker")

async def cpp_ws_loop() -> None:
    frames_total = 0
    skipped_total = 0
    last_log = time.time()

    while True:
        try:
            logger.info("[CPP WS] connecting to %s ...", CPP_WS_URI)
            async with websockets.connect(
                CPP_WS_URI, max_size=None, compression=None, ping_interval=None, close_timeout=1.0
            ) as ws:
                logger.info("[CPP WS] connected")
                async for raw in ws:
                    if isinstance(raw, (bytes, bytearray)):
                        buf = memoryview(raw)
                        if len(buf) < 3 or buf[0] != 1:
                            continue
                        idx = int.from_bytes(buf[1:3], "big")
                        from config import CAMERA_CONFIG
                        if idx not in CAMERA_CONFIG:
                            continue
                        if DROP_DECODE_WHEN_BUSY and ptz_states[idx]["proc_busy"] and (not video_clients):
                            skipped_total += 1
                            continue
                        arr = np.frombuffer(buf[3:], dtype=np.uint8)
                        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
                        if img is None:
                            continue
                    else:
                        try:
                            msg = json_loads(raw)
                        except Exception:
                            continue
                        if msg.get("type") != "image":
                            continue
                        idx_str, b64 = msg["data"].split("|", 1)
                        idx = int(idx_str)
                        from config import CAMERA_CONFIG
                        if idx not in CAMERA_CONFIG:
                            continue
                        if DROP_DECODE_WHEN_BUSY and ptz_states[idx]["proc_busy"] and (not video_clients):
                            skipped_total += 1
                            continue
                        arr = np.frombuffer(base64.b64decode(b64), dtype=np.uint8)
                        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
                        if img is None:
                            continue

                    frames_total += 1
                    h, w = img.shape[:2]
                    now = time.time()
                    if now - last_log >= 1.0:
                        logger.info("[CPP WS] recv: %d fps, skipped=%d | last cam=%s | %dx%d | clients=%d",
                                    frames_total, skipped_total, idx, w, h, len(video_clients))
                        frames_total = 0; skipped_total = 0; last_log = now

                    detected_cameras.add(idx)

                    if PUBLISH_RAW_BEFORE_INFER and video_clients:
                        down, _, _ = maybe_downscale(img, PREVIEW_TARGET_W)
                        publish_preview(idx, down)

                    pst = ptz_states[idx]
                    if pst["proc_busy"]:
                        continue
                    pst["proc_busy"] = True

                    async def process_frame(cam_idx: int, frame):
                        try:
                            assert batcher is not None, "batcher is not initialized"
                            await batcher.submit(cam_idx, frame, cam_idx in from_state_ptz_ids())
                        except Exception as exc:
                            logger.error("Error processing frame for camera %s: %s", cam_idx, exc)
                        finally:
                            ptz_states[cam_idx]["proc_busy"] = False

                    # helper to avoid circular import at top-level
                    def from_state_ptz_ids():
                        from state import PTZ_CAM_IDS
                        return PTZ_CAM_IDS

                    asyncio.create_task(process_frame(idx, img))
        except Exception as exc:
            logger.error("[CPP WS] error: %s", exc)
            await asyncio.sleep(0.5)

async def cpp_detection_loop() -> None:
    assert detection_queue is not None
    RESYNC_EVERY_SEC = 10.0

    def _snapshot_states():
        from config import CAMERA_CONFIG
        snap = []
        for cid in CAMERA_CONFIG.keys():
            st = ptz_states.get(cid, {})
            snap.append({"IdCamera": int(cid), "detection": bool(st.get("rec_active", False))})
        return snap

    while True:
        try:
            logger.info("[CPP CTRL] connecting to %s ...", CPP_CTRL_WS_URI)
            async with websockets.connect(
                CPP_CTRL_WS_URI, max_size=None, compression=None, ping_interval=None, close_timeout=1.0
            ) as ws:
                logger.info("[CPP CTRL] connected")
                try:
                    for item in _snapshot_states():
                        await ws.send(json_dumps({"type": "detection", "data": item}))
                except Exception as e:
                    logger.error("[CPP CTRL] initial resync failed: %s", e)
                    raise

                last_resync = time.time()

                while True:
                    now = time.time()
                    if now - last_resync >= RESYNC_EVERY_SEC and detection_queue.empty():
                        try:
                            for item in _snapshot_states():
                                await ws.send(json_dumps({"type": "detection", "data": item}))
                        except Exception as e:
                            logger.error("[CPP CTRL] periodic resync failed: %s", e)
                            break
                        last_resync = now

                    try:
                        payload = await asyncio.wait_for(detection_queue.get(), timeout=1.0)
                    except asyncio.TimeoutError:
                        continue

                    try:
                        await ws.send(json_dumps(payload))
                        last_resync = time.time()
                    except Exception as exc:
                        logger.error("[CPP CTRL] send failed: %s", exc)
                        try:
                            detection_queue.put_nowait(payload)
                        except Exception:
                            pass
                        break
                    finally:
                        detection_queue.task_done()

        except Exception as exc:
            logger.error("[CPP CTRL] error: %s", exc)
        await asyncio.sleep(1.0)