Enhance HVideoTool's detection and restoration capabilities: introduced per-model overlay threshold settings and cross-model non-maximum suppression (NMS) to improve detection accuracy. Updated configuration management to support these features, and refined the UI for better user experience. Documentation in CLAUDE.md has been updated to reflect these changes.
This commit is contained in:
@@ -15,10 +15,13 @@ from .types import Detection
|
||||
|
||||
|
||||
class MultiYoloDetector(Detector):
|
||||
def __init__(self, detectors: list[Detector]) -> None:
|
||||
def __init__(self, detectors: list[Detector], nms_iou: float | None = None) -> None:
|
||||
if not detectors:
|
||||
raise ValueError("MultiYoloDetector requires at least one detector")
|
||||
self._detectors = detectors
|
||||
# When set, overlapping detections (across all models, regardless of category)
|
||||
# are merged by greedy IoU NMS — the higher-score box wins. None = keep all.
|
||||
self._nms_iou = nms_iou
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -28,4 +31,33 @@ class MultiYoloDetector(Detector):
|
||||
out: list[Detection] = []
|
||||
for d in self._detectors:
|
||||
out.extend(d.detect(frame))
|
||||
if self._nms_iou is not None:
|
||||
out = _nms(out, self._nms_iou)
|
||||
return out
|
||||
|
||||
|
||||
def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
|
||||
"""Intersection-over-union of two (x, y, w, h) boxes."""
|
||||
ax, ay, aw, ah = a
|
||||
bx, by, bw, bh = b
|
||||
ix1, iy1 = max(ax, bx), max(ay, by)
|
||||
ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
|
||||
iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
|
||||
inter = iw * ih
|
||||
if inter == 0:
|
||||
return 0.0
|
||||
union = aw * ah + bw * bh - inter
|
||||
return inter / union if union > 0 else 0.0
|
||||
|
||||
|
||||
def _nms(dets: list[Detection], iou_thresh: float) -> list[Detection]:
|
||||
"""Greedy non-maximum suppression across all detections (category-agnostic).
|
||||
|
||||
Highest score first; a box is dropped if it overlaps an already-kept box by more
|
||||
than ``iou_thresh``. Used to remove duplicate boxes from overlapping models.
|
||||
"""
|
||||
kept: list[Detection] = []
|
||||
for d in sorted(dets, key=lambda x: x.score, reverse=True):
|
||||
if all(_iou(d.bbox, k.bbox) <= iou_thresh for k in kept):
|
||||
kept.append(d)
|
||||
return kept
|
||||
|
||||
Reference in New Issue
Block a user