114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
"""StereoYoloFrontend —— 双目检测前端主类（需求说明 §2/§3）。
|
||
|
||
单帧流程：左右图（已校正、原生分辨率）→ 左右并发检测 → 出口自检 + 类别
|
||
白名单 → 每目封装 DetectionFrame（共用 frame_id，§2.3 同帧契约）。
|
||
|
||
前端**无跨帧状态**：没有历史就没有外推，结构性满足"只送本帧原生实测框"
|
||
的输入纯净性硬约束（§2.5 第 1 条）。
|
||
"""
|
||
|
||
import logging
|
||
import math
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from typing import List, Optional, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
|
||
from stereo_ranger.types import Detection, DetectionFrame
|
||
|
||
from .config import StereoYoloConfig
|
||
from .detector import UltralyticsDetector
|
||
from .warn_throttle import WarnThrottle
|
||
|
||
logger = logging.getLogger("stereo_yolo")
|
||
|
||
_EYE_LABELS = ("左目", "右目")
|
||
|
||
|
||
class StereoYoloFrontend:
|
||
"""图像 → 符合 stereo_ranger 输入契约的检测框（§2 接口契约）。
|
||
|
||
``detector`` 可注入替换：接口为 ``detect(img) -> list[Detection]``
|
||
（单测用假检测器驱动，§6 输出合法性/失败隔离用例无需真实模型）。
|
||
"""
|
||
|
||
def __init__(self, config: StereoYoloConfig, detector: Optional[object] = None) -> None:
|
||
self._cfg = config
|
||
self._detector = detector if detector is not None else UltralyticsDetector(config)
|
||
self._throttle = WarnThrottle()
|
||
self._executor = (
|
||
ThreadPoolExecutor(max_workers=2, thread_name_prefix="stereo_yolo")
|
||
if config.concurrency == "parallel"
|
||
else None
|
||
)
|
||
|
||
@property
|
||
def config(self) -> StereoYoloConfig:
|
||
return self._cfg
|
||
|
||
def detect(
|
||
self,
|
||
left_img: np.ndarray,
|
||
right_img: np.ndarray,
|
||
frame_id: int,
|
||
) -> Tuple[DetectionFrame, DetectionFrame]:
|
||
"""左右各输出一个 DetectionFrame；两帧共用同一 frame_id（§2.3）。"""
|
||
if left_img.shape != right_img.shape:
|
||
self._warn(
|
||
"shape_mismatch",
|
||
f"左右图尺寸不一致（{left_img.shape} vs {right_img.shape}），整帧拒收输出空（§2.2 第 2 条）",
|
||
)
|
||
return DetectionFrame(frame_id, ()), DetectionFrame(frame_id, ())
|
||
|
||
if self._executor is not None:
|
||
left_future = self._executor.submit(self._detect_eye, left_img, _EYE_LABELS[0])
|
||
right_future = self._executor.submit(self._detect_eye, right_img, _EYE_LABELS[1])
|
||
left, right = left_future.result(), right_future.result()
|
||
else:
|
||
left = self._detect_eye(left_img, _EYE_LABELS[0])
|
||
right = self._detect_eye(right_img, _EYE_LABELS[1])
|
||
|
||
return DetectionFrame(frame_id, tuple(left)), DetectionFrame(frame_id, tuple(right))
|
||
|
||
def close(self) -> None:
|
||
"""释放并发线程池（前端不再使用时调用）。"""
|
||
if self._executor is not None:
|
||
self._executor.shutdown(wait=True)
|
||
|
||
# ---------------- 内部 ----------------
|
||
|
||
def _detect_eye(self, image: np.ndarray, eye_label: str) -> List[Detection]:
|
||
try:
|
||
raw = self._detector.detect(image)
|
||
except Exception as exc: # 单目失败隔离：不阻塞整体，该目输出空帧列表（§5 失败处理）
|
||
self._warn(f"infer_failed:{eye_label}", f"{eye_label}推理失败，该目输出空帧列表：{exc!r}")
|
||
return []
|
||
return self._sanitize(raw)
|
||
|
||
def _sanitize(self, detections: Sequence[Detection]) -> List[Detection]:
|
||
"""出口自检（§2.2 第 6 条）：丢弃非有限值 / x2<=x1 / y2<=y1 的框；白名单过滤。"""
|
||
kept: List[Detection] = []
|
||
dropped = 0
|
||
whitelist = self._cfg.class_whitelist
|
||
for detection in detections:
|
||
if not _is_valid_box(detection):
|
||
dropped += 1
|
||
continue
|
||
if whitelist is not None and detection.class_id not in whitelist:
|
||
continue
|
||
kept.append(detection)
|
||
if dropped:
|
||
self._warn("degenerate_box", f"丢弃退化检测框 {dropped} 个（非有限值或 x2<=x1/y2<=y1）")
|
||
return kept
|
||
|
||
def _warn(self, category: str, message: str) -> None:
|
||
if self._throttle.should_emit(category):
|
||
logger.warning("%s（累计 %d 次）", message, self._throttle.count(category))
|
||
|
||
|
||
def _is_valid_box(detection: Detection) -> bool:
|
||
values = (detection.x1, detection.y1, detection.x2, detection.y2)
|
||
if not all(math.isfinite(value) for value in values):
|
||
return False
|
||
return detection.x2 > detection.x1 and detection.y2 > detection.y1
|