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

if_curator/embeddings.py

@@ -1,8 +1,9 @@
 """
 Unified embedding interface for faces and objects.
 
-- Faces: InsightFace (ArcFace/Buffalo_L)
+- Faces: InsightFace (ArcFace/Buffalo_L) — or reuse from Immich
 - Objects: SigLIP (Vision Transformer via transformers)
+- Caching: Disk-based cache avoids recomputation on reruns
 """
 
 import contextlib
@@ -14,6 +15,8 @@ import cv2
 import numpy as np
 from PIL import Image
 
+from .cache import get_cache
+
 logger = logging.getLogger(__name__)
 
 # Lazy-loaded singletons
@@ -48,8 +51,10 @@ def get_insightface_app():
 
         # Determine device: 0 for GPU, -1 for CPU
         gpu_providers = {
-            "CUDAExecutionProvider", "ROCmExecutionProvider",
-            "MPSExecutionProvider", "CoreMLExecutionProvider",
+            "CUDAExecutionProvider",
+            "ROCmExecutionProvider",
+            "MPSExecutionProvider",
+            "CoreMLExecutionProvider",
         }
         ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1)
 
@@ -73,6 +78,7 @@ def get_insightface_app():
             logger.warning("Retrying InsightFace on CPU...")
             try:
                 from insightface.app import FaceAnalysis
+
                 _insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface")
                 _insightface_app.prepare(ctx_id=-1, det_size=(640, 640))
                 return _insightface_app
@@ -179,14 +185,83 @@ def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
         return None
 
 
+def get_object_embeddings_batch(images: list[Image.Image]) -> list[np.ndarray | None]:
+    """Get SigLIP embeddings for a batch of images (GPU-efficient)."""
+    model, processor = get_siglip_model()
+    if model is None:
+        return [None] * len(images)
+
+    try:
+        import torch
+
+        inputs = processor(images=images, return_tensors="pt", padding=True)
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+            embeddings = outputs.pooler_output.cpu().numpy()
+            return [embeddings[i] for i in range(len(embeddings))]
+    except Exception as e:
+        logger.error(f"Error in batch embedding: {e}")
+        # Fall back to individual computation
+        return [get_object_embedding(img) for img in images]
+
+
 # =============================================================================
-# Unified Interface
+# Unified Interface with Caching
 # =============================================================================
 
 
-def get_embedding(img_pil: Image.Image, entity_type: str = "face") -> np.ndarray | None:
-    """Get embedding for an image based on entity type ('face' or 'object')."""
-    return get_face_embedding(img_pil) if entity_type == "face" else get_object_embedding(img_pil)
+def get_embedding(
+    img_pil: Image.Image,
+    entity_type: str = "face",
+    asset_id: str | None = None,
+    immich_embedding: np.ndarray | None = None,
+) -> np.ndarray | None:
+    """Get embedding for an image based on entity type.
+
+    Priority:
+    1. Pre-fetched Immich embedding (if provided)
+    2. Disk cache (if enabled and asset_id provided)
+    3. Local model computation (InsightFace or SigLIP)
+
+    Args:
+        img_pil: The image to embed
+        entity_type: 'face' or 'object'
+        asset_id: Optional asset ID for cache lookup
+        immich_embedding: Optional pre-fetched embedding from Immich API
+    """
+    from .config import Config
+
+    use_cache = Config.ENABLE_CACHE and asset_id is not None
+    cache = get_cache(Config.CACHE_DIR) if use_cache else None
+    model_key = "immich" if entity_type == "face" else "siglip"
+
+    # 1. Use Immich embedding if provided
+    if immich_embedding is not None:
+        if cache:
+            cache.put(asset_id, immich_embedding, model_key)
+        return immich_embedding
+
+    # 2. Check disk cache
+    if cache:
+        cached = cache.get(asset_id, model_key)
+        if cached is not None:
+            return cached
+
+    # 3. Compute locally
+    if entity_type == "face":
+        emb = get_face_embedding(img_pil)
+        model_key = "insightface"
+    else:
+        emb = get_object_embedding(img_pil)
+
+    # Cache the result
+    if emb is not None and cache:
+        cache.put(asset_id, emb, model_key)
+
+    return emb
 
 
 def is_embedding_available(entity_type: str = "face") -> bool: