video_monitor/app/analysis/remote_detector.py

"""远程推理代理 — 摄像头子进程通过 Queue 向主进程 InferenceProcessPool 发起推理"""
import logging
import threading
import uuid

logger = logging.getLogger("analysis.remote_detector")

# 同一 resp_queue 只能有一个 drain 线程，否则多 RemoteDetector 会抢响应导致丢包
_drainers = {}
_drainers_lock = threading.Lock()


class _SharedResponseDrainer(object):
    def __init__(self, resp_queue):
        self._resp_q = resp_queue
        self._pending = {}
        self._lock = threading.Lock()
        self._started = False

    def ensure_started(self):
        if self._started:
            return
        self._started = True
        threading.Thread(
            target=self._loop, name="remote-det-drain-%s" % id(self._resp_q), daemon=True,
        ).start()

    def register(self, req_id):
        evt = {"event": threading.Event(), "resp": None}
        with self._lock:
            self._pending[req_id] = evt
        return evt

    def unregister(self, req_id):
        with self._lock:
            self._pending.pop(req_id, None)

    def _loop(self):
        import queue as _q
        while True:
            try:
                msg = self._resp_q.get(timeout=1.0)
            except _q.Empty:
                continue
            if not msg:
                continue
            req_id = msg.get("req_id")
            with self._lock:
                item = self._pending.pop(req_id, None)
            if item:
                item["resp"] = msg
                item["event"].set()
            elif req_id:
                logger.debug("remote_detector: 无匹配 pending req_id=%s", req_id)


def _get_drainer(resp_queue):
    key = id(resp_queue)
    with _drainers_lock:
        drainer = _drainers.get(key)
        if drainer is None:
            drainer = _SharedResponseDrainer(resp_queue)
            _drainers[key] = drainer
        return drainer


class RemoteDetector(object):
    ENGINE_NAME = "remote_pool"

    def __init__(self, algorithm_spec, req_queue, resp_queue, timeout=35.0,
                 response_channel=None):
        self._spec = algorithm_spec
        self._req_q = req_queue
        self._resp_q = resp_queue
        self._timeout = timeout
        self._response_channel = response_channel
        self._drainer = _get_drainer(resp_queue)

    def ready(self):
        return self._req_q is not None and self._resp_q is not None

    def load(self):
        return True

    def detect(self, frame):
        if not self.ready():
            raise RuntimeError("共享推理队列未就绪")
        self._drainer.ensure_started()
        try:
            import cv2
            orig_h, orig_w = frame.shape[:2]
            inference_frame = self._maybe_downscale(frame)
            infer_h, infer_w = inference_frame.shape[:2]
            # 远处人员框很小，低质量 JPEG 会在 YOLO 前抹掉轮廓；95 在当前单路 CPU
            # 场景仍可控，并显著缩小与直接单图推理的差异。
            ok, buf = cv2.imencode(".jpg", inference_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
            if not ok:
                raise RuntimeError("推理帧编码失败")
            jpeg = buf.tobytes()
        except Exception as e:
            raise RuntimeError("推理帧编码失败: %s" % e) from e

        req_id = str(uuid.uuid4())
        evt = self._drainer.register(req_id)
        try:
            self._req_q.put({
                "req_id": req_id,
                "algorithm": self._spec,
                "jpeg": jpeg,
                "response_channel": self._response_channel,
            }, timeout=2.0)
        except Exception as e:
            self._drainer.unregister(req_id)
            raise RuntimeError("推理请求入队失败: %s" % e) from e

        if not evt["event"].wait(timeout=self._timeout):
            self._drainer.unregister(req_id)
            logger.warning("RemoteDetector 推理超时 algo=%s", self._spec.get("name"))
            raise TimeoutError("共享推理响应超时（%ss）" % self._timeout)
        resp = evt.get("resp") or {}
        if not resp.get("ok"):
            raise RuntimeError(resp.get("error") or "共享推理失败")
        detections = resp.get("detections") or []
        if infer_w != orig_w or infer_h != orig_h:
            detections = self._restore_coordinates(
                detections, float(orig_w) / infer_w, float(orig_h) / infer_h,
            )
        return detections

    @staticmethod
    def _restore_coordinates(detections, scale_x, scale_y):
        """把共享推理缩图坐标恢复到解码原帧坐标系。"""
        result = []
        for source in detections or []:
            item = dict(source)
            box = item.get("box")
            if isinstance(box, (list, tuple)) and len(box) >= 4:
                item["box"] = [float(box[0]) * scale_x, float(box[1]) * scale_y,
                               float(box[2]) * scale_x, float(box[3]) * scale_y]
            keypoints = item.get("keypoints")
            if isinstance(keypoints, list):
                item["keypoints"] = [
                    [float(p[0]) * scale_x, float(p[1]) * scale_y] + list(p[2:])
                    if isinstance(p, (list, tuple)) and len(p) >= 2 else p
                    for p in keypoints
                ]
            for polygon_key in ("polygon", "mask", "segments"):
                polygon = item.get(polygon_key)
                if isinstance(polygon, list):
                    item[polygon_key] = RemoteDetector._scale_polygon(polygon, scale_x, scale_y)
            result.append(item)
        return result

    @staticmethod
    def _scale_polygon(value, scale_x, scale_y):
        scaled = []
        for point in value:
            if (isinstance(point, (list, tuple)) and len(point) >= 2 and
                    isinstance(point[0], (int, float)) and isinstance(point[1], (int, float))):
                scaled.append([float(point[0]) * scale_x, float(point[1]) * scale_y] + list(point[2:]))
            elif isinstance(point, list):
                scaled.append(RemoteDetector._scale_polygon(point, scale_x, scale_y))
            else:
                scaled.append(point)
        return scaled

    def _maybe_downscale(self, frame):
        try:
            import cv2
            h, w = frame.shape[:2]
            max_side = max(
                int(self._spec.get("input_width", 640) or 640),
                int(self._spec.get("input_height", 640) or 640),
                640,
            ) * 4
            longest = max(h, w)
            if longest <= max_side:
                return frame
            scale = float(max_side) / float(longest)
            nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
            return cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_AREA)
        except Exception:
            return frame