yolo

target_roi_yolo.py

@@ -89,6 +89,29 @@ def _stage2_roi_crop_save_worker(
 _detector_by_path = {}
 
 
+def _resolve_model_path(model_path: str):
+    """Resolve a model in either the installed app or MaixVision run directory."""
+    model_path = (model_path or "").strip()
+    if model_path and os.path.isfile(model_path):
+        return model_path
+    if not model_path:
+        return ""
+    name = os.path.basename(model_path)
+    module_dir = os.path.dirname(os.path.abspath(__file__))
+    candidates = (
+        os.path.join(module_dir, name),
+        os.path.join(module_dir, "test", name),
+        os.path.join("/tmp/maixpy_run", name),
+        os.path.join("/tmp/maixpy_run", "test", name),
+        os.path.join(os.getcwd(), name),
+        os.path.join(os.getcwd(), "test", name),
+    )
+    for candidate in candidates:
+        if os.path.isfile(candidate):
+            return candidate
+    return model_path
+
+
 def reset_yolo_detector_cache():
     """切换模型路径时可调用（通常不必）。"""
     global _detector_by_path
@@ -175,6 +198,23 @@ def preload_yolo_detector(logger=None):
                 % (_loc_black,)
             )
 
+    if bool(getattr(cfg, "TARGET_CLASS_YOLO_ENABLE", False)) and bool(
+        getattr(cfg, "TARGET_CLASS_YOLO_PRELOAD_ON_BOOT", True)
+    ):
+        class_model_path = _resolve_model_path(
+            getattr(cfg, "TARGET_CLASS_YOLO_MODEL_PATH", "") or ""
+        )
+        class_detector = _get_detector(class_model_path)
+        if class_detector is None:
+            if logger:
+                logger.warning(
+                    f"[YOLO-TARGET] 预加载失败：无法加载模型 {class_model_path}"
+                )
+        else:
+            ok = True
+            if logger:
+                logger.info(f"[YOLO-TARGET] 靶规格模型已预加载: {class_model_path}")
+
     return ok
 
 
@@ -206,8 +246,10 @@ def _det_obj_class_id(o):
             if v is None:
                 continue
             try:
+                if callable(v):
+                    v = v()
                 return int(float(v))
-            except (TypeError, ValueError):
+            except (TypeError, ValueError, AttributeError):
                 continue
     return None
 
@@ -242,6 +284,106 @@ def _normalize_objs(objs):
     return out
 
 
+def _det_obj_score(o):
+    """兼容 Maix YOLO 不同版本的置信度字段。"""
+    for key in ("score", "confidence", "conf", "prob"):
+        if hasattr(o, key):
+            try:
+                value = getattr(o, key)
+                if callable(value):
+                    value = value()
+                value = float(value)
+                if value == value:
+                    return value
+            except (TypeError, ValueError, AttributeError):
+                pass
+    return 0.0
+
+
+def try_get_target_class_from_yolo(maix_frame, logger=None):
+    """识别当前帧的 20/40 靶规格，失败返回 None。"""
+    try:
+        import config as cfg
+    except Exception:
+        return None
+
+    if not bool(getattr(cfg, "TARGET_CLASS_YOLO_ENABLE", False)):
+        return None
+    model_path = _resolve_model_path(
+        getattr(cfg, "TARGET_CLASS_YOLO_MODEL_PATH", "") or ""
+    )
+    if not os.path.isfile(model_path):
+        if logger:
+            logger.warning(f"[YOLO-TARGET] 模型文件不存在: {model_path}")
+        return None
+    detector = _get_detector(model_path)
+    if detector is None:
+        if logger:
+            logger.warning("[YOLO-TARGET] 无法加载 nn.YOLOv5")
+        return None
+
+    conf_th = float(getattr(cfg, "TARGET_CLASS_YOLO_CONF_TH", 0.5))
+    iou_th = float(getattr(cfg, "TARGET_CLASS_YOLO_IOU_TH", 0.45))
+    labels = getattr(cfg, "TARGET_CLASS_YOLO_LABELS", (20, 40))
+    if isinstance(labels, str):
+        labels = tuple(x.strip() for x in labels.split(",") if x.strip())
+    labels = tuple(labels)
+
+    def _detect(threshold):
+        try:
+            raw = detector.detect(maix_frame, conf_th=threshold, iou_th=iou_th)
+        except Exception as exc:
+            if logger:
+                logger.warning(f"[YOLO-TARGET] detect 异常: {exc}")
+            return []
+        return _normalize_objs(raw if raw is not None else [])
+
+    def _candidates(objs):
+        found = []
+        for obj in objs:
+            class_id = _det_obj_class_id(obj)
+            if class_id is None or class_id < 0 or class_id >= len(labels):
+                continue
+            try:
+                label = int(float(labels[class_id]))
+            except (TypeError, ValueError):
+                continue
+            if label in (20, 40):
+                found.append((label, class_id, _det_obj_score(obj)))
+        return found
+
+    objects = _detect(conf_th)
+    candidates = _candidates(objects)
+    if logger and objects:
+        logger.info(
+            "[YOLO-TARGET] 原始框=%d, 解析类别=%s"
+            % (
+                len(objects),
+                [(_det_obj_class_id(o), _det_obj_score(o)) for o in objects[:8]],
+            )
+        )
+    if not candidates and bool(
+        getattr(cfg, "TARGET_CLASS_YOLO_RETRY_ON_EMPTY", False)
+    ):
+        retry_th = float(getattr(cfg, "TARGET_CLASS_YOLO_RETRY_CONF_TH", conf_th))
+        if 0 < retry_th < conf_th:
+            candidates = _candidates(_detect(retry_th))
+
+    if not candidates:
+        if logger:
+            logger.warning("[YOLO-TARGET] 当前帧未识别到 20/40，按未知处理")
+        return None
+
+    label, class_id, confidence = max(candidates, key=lambda item: item[2])
+    result = {"label": label, "class_id": class_id, "confidence": confidence}
+    if logger:
+        logger.info(
+            f"[YOLO-TARGET] 当前帧分类={label}, class_id={class_id}, "
+            f"conf={confidence:.3f}"
+        )
+    return result
+
+
 def _det_to_src_xyxy(o, coord_mode: str, src_w: int, src_h: int, net_w: int, net_h: int):
     """把单个检测框转为全图坐标系下的 xyxy（半开区间语义与后续 clip 一致）。"""
     x, y, w, h = float(o.x), float(o.y), float(o.w), float(o.h)