feat: Implement advanced diversity selection algorithms, comprehensive quality filtering, and caching for improved image curation.

if_curator/cache.py

@@ -0,0 +1,87 @@
+"""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
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Model versions — bump these when the upstream model changes
+MODEL_VERSIONS = {
+    "insightface": "buffalo_l_v1",
+    "siglip": "siglip-base-patch16-224_v1",
+    "immich": "immich_buffalo_l_v1",
+}
+
+
+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).
+    """
+
+    def __init__(self, cache_dir: str = ".if_cache") -> None:
+        self.cache_dir = cache_dir
+        self._ensured = False
+
+    def _ensure_dir(self) -> None:
+        if not self._ensured:
+            os.makedirs(self.cache_dir, exist_ok=True)
+            self._ensured = True
+
+    @staticmethod
+    def _key(asset_id: str, model: str) -> str:
+        version = 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 Exception:
+                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()
+        try:
+            np.save(self._path(asset_id, model), embedding)
+        except Exception as e:
+            logger.debug(f"Cache write failed for {asset_id}: {e}")
+
+    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"):
+                os.remove(os.path.join(self.cache_dir, f))
+                count += 1
+        logger.info(f"Cleared {count} cached embeddings.")
+
+
+# Singleton instance
+_cache: EmbeddingCache | None = None
+
+
+def get_cache(cache_dir: str = ".if_cache") -> EmbeddingCache:
+    """Get or create the singleton cache instance."""
+    global _cache
+    if _cache is None:
+        _cache = EmbeddingCache(cache_dir)
+    return _cache