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:
@@ -17,6 +17,7 @@ project folder.
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from .types import Detection
|
||||
@@ -24,16 +25,43 @@ from .types import Detection
|
||||
_VERSION = 1
|
||||
|
||||
|
||||
def make_key(models: list[str], yolo_conf: float, yolo_imgsz: int) -> dict:
|
||||
def _atomic_write_text(path: Path, text: str) -> None:
|
||||
"""Write ``text`` to ``path`` crash-safely: write a sibling .tmp, then os.replace.
|
||||
|
||||
``os.replace`` is atomic on the same filesystem (incl. NTFS), so a crash mid-write
|
||||
leaves the previous file intact instead of a truncated/corrupt one — important for
|
||||
a large detections.json that holds tens of thousands of entries.
|
||||
"""
|
||||
tmp = path.with_name(path.name + ".tmp")
|
||||
try:
|
||||
tmp.write_text(text, encoding="utf-8")
|
||||
os.replace(tmp, path)
|
||||
finally:
|
||||
if tmp.exists():
|
||||
try:
|
||||
tmp.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def make_key(
|
||||
models: list[str], yolo_conf: float, yolo_imgsz: int, nms_iou: float | None = None
|
||||
) -> dict:
|
||||
"""Identity of the detector set that produced a cache; cache is only valid for a match.
|
||||
|
||||
Keyed by the (sorted) model **basenames** so it's portable across machines/paths.
|
||||
``nms_iou`` is only added to the key when cross-model NMS is enabled — so the default
|
||||
(NMS off) key is unchanged and existing caches stay valid; turning NMS on yields a
|
||||
distinct key (its merged results differ) without invalidating the non-NMS cache.
|
||||
"""
|
||||
return {
|
||||
key = {
|
||||
"models": sorted(Path(m).name for m in models),
|
||||
"yolo_conf": round(float(yolo_conf), 4),
|
||||
"yolo_imgsz": int(yolo_imgsz),
|
||||
}
|
||||
if nms_iou is not None:
|
||||
key["nms_iou"] = round(float(nms_iou), 4)
|
||||
return key
|
||||
|
||||
|
||||
def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]]) -> bool:
|
||||
@@ -48,9 +76,7 @@ def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]
|
||||
}
|
||||
try:
|
||||
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
cache_file.write_text(
|
||||
json.dumps(payload, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
_atomic_write_text(cache_file, json.dumps(payload, ensure_ascii=False))
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
@@ -31,6 +31,8 @@ def build_detector(config: AppConfig) -> Detector:
|
||||
from .multi import MultiYoloDetector
|
||||
from .yolo import YoloDetector
|
||||
|
||||
return MultiYoloDetector([
|
||||
YoloDetector(m, config.detection, label=category_of(m)) for m in models
|
||||
])
|
||||
nms_iou = config.nms_iou if config.cross_model_nms else None
|
||||
return MultiYoloDetector(
|
||||
[YoloDetector(m, config.detection, label=category_of(m)) for m in models],
|
||||
nms_iou=nms_iou,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -19,6 +19,7 @@ threshold, restore engine). Global ``settings.json`` only seeds the defaults for
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
@@ -38,6 +39,9 @@ _SETTING_KEYS = (
|
||||
"detector",
|
||||
"detector_models",
|
||||
"default_threshold",
|
||||
"model_thresholds",
|
||||
"cross_model_nms",
|
||||
"nms_iou",
|
||||
"restorer",
|
||||
"dm_dir",
|
||||
"dm_model",
|
||||
@@ -139,9 +143,19 @@ class Project:
|
||||
"settings": self.settings,
|
||||
}
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self.project_file.write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
# Crash-safe write (sibling .tmp + atomic os.replace): never truncate a good
|
||||
# project.json if the app/PC dies mid-save.
|
||||
text = json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
tmp = self.project_file.with_name(PROJECT_FILE + ".tmp")
|
||||
try:
|
||||
tmp.write_text(text, encoding="utf-8")
|
||||
os.replace(tmp, self.project_file)
|
||||
finally:
|
||||
if tmp.exists():
|
||||
try:
|
||||
tmp.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
# ----------------------------------------------------- settings <-> config
|
||||
def apply_to_config(self, cfg: AppConfig) -> None:
|
||||
|
||||
Reference in New Issue
Block a user