flan
4af9da2550
- diversity: scale face bbox to thumbnail space before quality check so
check_face_size uses actual thumbnail pixels, not original-image coords
- diversity: skip asset when face bbox exists but crop guard rejects it,
preventing InsightFace from picking the wrong person in a group photo
- diversity: add _scale_bbox_to_thumbnail helper (extracted from crop logic)
- diversity: use set for medoid membership test in _kmedoids (O(n) not O(n*k))
- diversity: remove dead np.unique in _select_time_spread (linspace produces
strictly increasing indices; unique is a no-op and implies wrong semantics)
- embeddings: move os.open/os.dup calls inside try in _suppress_output so
EMFILE during setup does not leak already-allocated fds
- immich_api: count and log assets with missing/unparseable fileCreatedAt in
filter_recent_assets instead of silently discarding them
- executor: capture pre_run_count before stale-mapping cleanup so the
"first run" coaching message doesn't fire after manual file deletion
- cli: use p['id'] (KeyError-safe) instead of p.get('id') in fallback path
to match all other access sites on the same people list
- cache: narrow except to (OSError, ValueError) in EmbeddingCache.get so
MemoryError propagates instead of converting OOM to a silent cache miss
131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
"""Disk-based embedding cache.
|
|
|
|
Caches embeddings keyed by (asset_id, model_version) to avoid
|
|
recomputing on reruns. Uses numpy binary format for fast I/O.
|
|
"""
|
|
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MODEL_VERSIONS = {
|
|
"immich": "immich_buffalo_l_v1",
|
|
}
|
|
|
|
|
|
def _insightface_model_fingerprint() -> str:
|
|
"""Derive a version string from buffalo_l .onnx file sizes and mtimes.
|
|
|
|
Changes automatically when model files are replaced or updated, preventing
|
|
stale embeddings from a previous model being served from cache.
|
|
Falls back to a static string before the model is downloaded (first run).
|
|
"""
|
|
insightface_home = os.environ.get("INSIGHTFACE_HOME", os.path.expanduser("~/.insightface"))
|
|
model_dir = Path(insightface_home) / "models" / "buffalo_l"
|
|
if not model_dir.exists():
|
|
return "buffalo_l_v1"
|
|
onnx_files = sorted(model_dir.glob("*.onnx"))
|
|
if not onnx_files:
|
|
return "buffalo_l_v1"
|
|
fingerprint = "|".join(
|
|
f"{f.name}:{f.stat().st_size}:{int(f.stat().st_mtime)}"
|
|
for f in onnx_files
|
|
)
|
|
return hashlib.sha256(fingerprint.encode()).hexdigest()[:12]
|
|
|
|
|
|
class EmbeddingCache:
|
|
"""Simple disk-based embedding cache.
|
|
|
|
Embeddings are stored as .npy files in a flat directory,
|
|
keyed by a hash of (asset_id, model_version). The InsightFace version
|
|
is derived from buffalo_l model file metadata so the cache auto-invalidates
|
|
when model files are replaced or updated.
|
|
"""
|
|
|
|
def __init__(self, cache_dir: str = ".if_cache") -> None:
|
|
self.cache_dir = cache_dir
|
|
self._ensured = False
|
|
self._model_versions = {
|
|
**MODEL_VERSIONS,
|
|
"insightface": _insightface_model_fingerprint(),
|
|
}
|
|
|
|
def _ensure_dir(self) -> None:
|
|
if not self._ensured:
|
|
os.makedirs(self.cache_dir, exist_ok=True)
|
|
self._ensured = True
|
|
|
|
def _key(self, asset_id: str, model: str) -> str:
|
|
version = self._model_versions.get(model, model)
|
|
raw = f"{asset_id}:{version}"
|
|
return hashlib.sha256(raw.encode()).hexdigest()[:16]
|
|
|
|
def _path(self, asset_id: str, model: str) -> str:
|
|
return os.path.join(self.cache_dir, f"{self._key(asset_id, model)}.npy")
|
|
|
|
def get(self, asset_id: str, model: str = "insightface") -> np.ndarray | None:
|
|
"""Retrieve cached embedding, or None if not cached."""
|
|
path = self._path(asset_id, model)
|
|
if os.path.exists(path):
|
|
try:
|
|
return np.load(path)
|
|
except (OSError, ValueError):
|
|
return None
|
|
return None
|
|
|
|
def put(self, asset_id: str, embedding: np.ndarray, model: str = "insightface") -> None:
|
|
"""Store an embedding in the cache."""
|
|
self._ensure_dir()
|
|
final = self._path(asset_id, model)
|
|
# Insert .tmp before .npy so np.save doesn't auto-append another .npy extension
|
|
# (np.save appends .npy to paths that don't already end in .npy).
|
|
tmp = final.removesuffix(".npy") + ".tmp.npy"
|
|
try:
|
|
np.save(tmp, embedding)
|
|
os.replace(tmp, final)
|
|
except Exception as e:
|
|
logger.warning("Cache write failed for %s: %s", asset_id, e)
|
|
try:
|
|
os.remove(tmp)
|
|
except OSError:
|
|
pass
|
|
|
|
def clear(self) -> None:
|
|
"""Delete all cached embeddings."""
|
|
if not os.path.isdir(self.cache_dir):
|
|
return
|
|
count = 0
|
|
for f in os.listdir(self.cache_dir):
|
|
if f.endswith(".npy"):
|
|
try:
|
|
os.remove(os.path.join(self.cache_dir, f))
|
|
count += 1
|
|
except OSError:
|
|
pass
|
|
logger.info("Cleared %s cached embeddings.", count)
|
|
|
|
|
|
# Singleton instance
|
|
_cache: EmbeddingCache | None = None
|
|
_cache_dir: str | None = None
|
|
|
|
|
|
def get_cache(cache_dir: str = ".if_cache") -> EmbeddingCache:
|
|
"""Get or create the singleton cache instance.
|
|
|
|
Re-creates the instance when ``cache_dir`` changes so that test
|
|
isolation (which resets Config.DATA_DIR via _Config.reset()) always
|
|
writes to the correct directory rather than a stale one.
|
|
"""
|
|
global _cache, _cache_dir
|
|
if _cache is None or _cache_dir != cache_dir:
|
|
_cache = EmbeddingCache(cache_dir)
|
|
_cache_dir = cache_dir
|
|
return _cache
|