79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
"""UltralyticsDetector —— ultralytics YOLO .pt 直载的检测后端。
|
||
|
||
⚠ 移植注记（C++ 侧需自行复现，详见 README）：letterbox（缩放 + 填充到
|
||
imgsz）、输出张量解码、NMS、框映射回输入图原生坐标系，本例由 ultralytics
|
||
库内完成；移植时必须逐项复现同一几何（步长对齐、填充方式、IoU 门限、
|
||
原生坐标回映），否则视差取点会引入系统性偏差。
|
||
"""
|
||
|
||
import queue
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from stereo_ranger.types import Detection
|
||
|
||
from .config import StereoYoloConfig
|
||
|
||
|
||
class UltralyticsDetector:
|
||
"""双实例池：ultralytics 单实例非线程安全，左右并发各持一份（§3/§4）。
|
||
|
||
接口约定（可注入替换）：``detect(image) -> list[Detection]``，
|
||
image 为 BGR ndarray（cv2 约定），返回框位于该图原生坐标系。
|
||
"""
|
||
|
||
def __init__(self, config: StereoYoloConfig) -> None:
|
||
self._cfg = config
|
||
self._model_path = _resolve_model_path(config.model_path)
|
||
self._pool: "queue.SimpleQueue" = queue.SimpleQueue()
|
||
from ultralytics import YOLO # 延迟导入：仅真实推理付加载成本，单测注入假检测器（§6）
|
||
|
||
for _ in range(2):
|
||
self._pool.put(YOLO(str(self._model_path)))
|
||
|
||
def detect(self, image) -> List[Detection]:
|
||
model = self._pool.get()
|
||
try:
|
||
results = model.predict(
|
||
source=image,
|
||
imgsz=self._cfg.imgsz,
|
||
conf=self._cfg.conf_threshold,
|
||
iou=self._cfg.iou_threshold,
|
||
device=self._cfg.device,
|
||
verbose=False,
|
||
)
|
||
finally:
|
||
self._pool.put(model)
|
||
if not results:
|
||
return []
|
||
return _result_to_detections(results[0])
|
||
|
||
|
||
def _result_to_detections(result) -> List[Detection]:
|
||
boxes = getattr(result, "boxes", None)
|
||
if boxes is None or len(boxes) == 0:
|
||
return []
|
||
xyxy = boxes.xyxy.cpu().numpy()
|
||
conf = boxes.conf.cpu().numpy()
|
||
cls = boxes.cls.cpu().numpy()
|
||
return [
|
||
Detection(
|
||
float(box[0]),
|
||
float(box[1]),
|
||
float(box[2]),
|
||
float(box[3]),
|
||
confidence=float(score),
|
||
class_id=int(class_id),
|
||
)
|
||
for box, score, class_id in zip(xyxy, conf, cls)
|
||
]
|
||
|
||
|
||
def _resolve_model_path(raw: str) -> Path:
|
||
path = Path(raw).expanduser()
|
||
if not path.is_absolute():
|
||
path = Path.cwd() / path
|
||
if not path.is_file():
|
||
raise FileNotFoundError(f"模型权重不存在: {path}（model_path 相对路径按启动工作目录解析）")
|
||
return path.resolve()
|