106 lines
3.8 KiB
Python
106 lines
3.8 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
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from .types import Detection
|
|
|
|
_VERSION = 1
|
|
|
|
|
|
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.
|
|
"""
|
|
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:
|
|
"""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)
|
|
_atomic_write_text(cache_file, json.dumps(payload, ensure_ascii=False))
|
|
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
|