52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
"""Composite detector: runs several detectors and merges their results.
|
|
|
|
Used for the "combined" mode = YOLO (mosaic) + classic-CV (black bars / blur).
|
|
Detections from all sub-detectors are concatenated, then de-duplicated by IoU
|
|
(higher score wins) so overlapping hits from different detectors don't stack.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from ..video.frame import Frame
|
|
from .base import Detector
|
|
from .types import Detection
|
|
|
|
|
|
class CompositeDetector(Detector):
|
|
def __init__(self, detectors: list[Detector], iou_threshold: float = 0.6) -> None:
|
|
if not detectors:
|
|
raise ValueError("CompositeDetector requires at least one detector")
|
|
self._detectors = detectors
|
|
self._iou = iou_threshold
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "Composite(" + " + ".join(d.name for d in self._detectors) + ")"
|
|
|
|
def detect(self, frame: Frame) -> list[Detection]:
|
|
merged: list[Detection] = []
|
|
for detector in self._detectors:
|
|
try:
|
|
merged += detector.detect(frame)
|
|
except Exception: # noqa: BLE001 - one detector failing must not kill the frame
|
|
continue
|
|
return self._dedup(merged)
|
|
|
|
def _dedup(self, dets: list[Detection]) -> list[Detection]:
|
|
dets = sorted(dets, key=lambda d: d.score, reverse=True)
|
|
kept: list[Detection] = []
|
|
for d in dets:
|
|
if all(self._iou_of(d.bbox, k.bbox) < self._iou for k in kept):
|
|
kept.append(d)
|
|
return kept
|
|
|
|
@staticmethod
|
|
def _iou_of(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
|
|
ax, ay, aw, ah = a
|
|
bx, by, bw, bh = b
|
|
ix, iy = max(ax, bx), max(ay, by)
|
|
ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
|
|
inter = max(0, ix2 - ix) * max(0, iy2 - iy)
|
|
union = aw * ah + bw * bh - inter
|
|
return inter / union if union > 0 else 0.0
|