78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
"""Persist detection results for a project.
|
|
|
|
The detection cache is a JSON file (``detections.json`` at the project root) that
|
|
maps each image to its detections so reopening a project doesn't have to re-run
|
|
the detector. The cache file lives apart from the images (which sit in the
|
|
project's ``frames/`` sub-folder), so the file location and the image base
|
|
directory are passed separately.
|
|
|
|
The cache is tagged with the detector identity (name + model + conf/imgsz); a
|
|
mismatch means the cache was produced by a different detector and is ignored
|
|
(``load_results`` returns ``None``) rather than shown as if current.
|
|
|
|
Keys are stored as **basenames**, so the cache survives moving/renaming the
|
|
project folder.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from .types import Detection
|
|
|
|
_VERSION = 1
|
|
|
|
|
|
def make_key(detector: str, model_path: str | None, yolo_conf: float, yolo_imgsz: int) -> dict:
|
|
"""Identity of the detector that produced a cache; cache is only valid for a match."""
|
|
return {
|
|
"detector": detector,
|
|
"model_path": model_path or "",
|
|
"yolo_conf": round(float(yolo_conf), 4),
|
|
"yolo_imgsz": int(yolo_imgsz),
|
|
}
|
|
|
|
|
|
def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]]) -> bool:
|
|
"""Write the cache (basename -> detections) to ``cache_file``. False on failure."""
|
|
payload = {
|
|
"version": _VERSION,
|
|
"key": key,
|
|
"results": {
|
|
Path(p).name: [d.to_dict() for d in dets]
|
|
for p, dets in results.items()
|
|
},
|
|
}
|
|
try:
|
|
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
|
cache_file.write_text(
|
|
json.dumps(payload, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def load_results(cache_file: Path, key: dict, base_dir: Path) -> dict[str, list[Detection]] | None:
|
|
"""Load cached detections from ``cache_file`` if present and the detector matches.
|
|
|
|
Returns a dict keyed by **full path** (``base_dir / basename``), or ``None`` if
|
|
there is no cache, it's unreadable, or it was made by a different detector.
|
|
"""
|
|
if not cache_file.is_file():
|
|
return None
|
|
try:
|
|
payload = json.loads(cache_file.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
return None
|
|
if payload.get("version") != _VERSION or payload.get("key") != key:
|
|
return None
|
|
out: dict[str, list[Detection]] = {}
|
|
for name, dets in payload.get("results", {}).items():
|
|
try:
|
|
out[str(base_dir / name)] = [Detection.from_dict(d) for d in dets]
|
|
except (KeyError, TypeError, ValueError):
|
|
continue # skip a corrupt entry, keep the rest
|
|
return out
|