273 lines
9.0 KiB
Python
273 lines
9.0 KiB
Python
"""
|
|
Unified embedding interface for faces and objects.
|
|
|
|
- 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
|
|
import logging
|
|
import os
|
|
import warnings
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from .cache import get_cache
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Lazy-loaded singletons
|
|
_insightface_app = None
|
|
_siglip_model = None
|
|
_siglip_processor = None
|
|
|
|
|
|
def _is_force_cpu() -> bool:
|
|
"""Check if CPU mode is forced via environment variable."""
|
|
return os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes")
|
|
|
|
|
|
# =============================================================================
|
|
# InsightFace (Faces)
|
|
# =============================================================================
|
|
|
|
|
|
def get_insightface_app():
|
|
"""Singleton for InsightFace app with automatic GPU/CPU fallback."""
|
|
global _insightface_app
|
|
if _insightface_app is not None:
|
|
return _insightface_app
|
|
|
|
try:
|
|
import onnxruntime as ort
|
|
from insightface.app import FaceAnalysis
|
|
|
|
# Get providers, excluding TensorRT to avoid noisy errors
|
|
providers = [p for p in ort.get_available_providers() if p != "TensorrtExecutionProvider"]
|
|
logger.info(f"Available ONNX providers: {providers}")
|
|
|
|
# Determine device: 0 for GPU, -1 for CPU
|
|
gpu_providers = {
|
|
"CUDAExecutionProvider",
|
|
"ROCmExecutionProvider",
|
|
"MPSExecutionProvider",
|
|
"CoreMLExecutionProvider",
|
|
}
|
|
ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1)
|
|
|
|
device_str = "GPU" if ctx_id >= 0 else "CPU"
|
|
logger.info(f"Loading InsightFace Buffalo_L on {device_str} (ctx_id={ctx_id})...")
|
|
|
|
# Suppress C-level output during model loading
|
|
with open(os.devnull, "w") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
|
|
_insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface", providers=providers)
|
|
_insightface_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
|
|
|
|
return _insightface_app
|
|
|
|
except ImportError:
|
|
logger.error("InsightFace not installed!")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Failed to load InsightFace: {e}")
|
|
# Retry on CPU if GPU failed
|
|
if ctx_id == 0:
|
|
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
|
|
except Exception as ex:
|
|
logger.error(f"CPU fallback failed: {ex}")
|
|
return None
|
|
|
|
|
|
def get_face_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
|
"""Get embedding of the largest face in a PIL image."""
|
|
app = get_insightface_app()
|
|
if not app:
|
|
return None
|
|
|
|
try:
|
|
# InsightFace expects BGR cv2 image
|
|
img_bgr = cv2.cvtColor(np.asarray(img_pil), cv2.COLOR_RGB2BGR)
|
|
|
|
# Suppress scikit-image FutureWarning from InsightFace's face_align.py
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*estimate.*is deprecated", category=FutureWarning)
|
|
faces = app.get(img_bgr)
|
|
|
|
if not faces:
|
|
return None
|
|
|
|
# Return embedding of largest face
|
|
largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
|
return largest.embedding
|
|
except Exception as e:
|
|
logger.error(f"Error getting face embedding: {e}")
|
|
return None
|
|
|
|
|
|
# =============================================================================
|
|
# SigLIP (Objects)
|
|
# =============================================================================
|
|
|
|
|
|
def get_siglip_model():
|
|
"""Singleton for SigLIP model and processor with GPU auto-detection."""
|
|
global _siglip_model, _siglip_processor
|
|
if _siglip_model is not None:
|
|
return _siglip_model, _siglip_processor
|
|
|
|
try:
|
|
import warnings
|
|
|
|
import torch
|
|
from transformers import AutoImageProcessor, SiglipVisionModel
|
|
|
|
model_name = "google/siglip-base-patch16-224"
|
|
logger.info(f"Loading SigLIP model ({model_name})...")
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", category=FutureWarning)
|
|
warnings.filterwarnings("ignore", message=".*use_fast.*")
|
|
_siglip_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
|
|
_siglip_model = SiglipVisionModel.from_pretrained(model_name)
|
|
|
|
_siglip_model.eval()
|
|
|
|
# Move to GPU if available
|
|
if not _is_force_cpu():
|
|
if torch.cuda.is_available():
|
|
_siglip_model = _siglip_model.cuda()
|
|
logger.info("SigLIP running on CUDA GPU")
|
|
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
_siglip_model = _siglip_model.to("mps")
|
|
logger.info("SigLIP running on Apple MPS")
|
|
else:
|
|
logger.info("SigLIP running on CPU")
|
|
else:
|
|
logger.info("FORCE_CPU set. SigLIP running on CPU")
|
|
|
|
return _siglip_model, _siglip_processor
|
|
|
|
except ImportError as e:
|
|
logger.error(f"transformers/torch not installed: {e}")
|
|
return None, None
|
|
except Exception as e:
|
|
logger.error(f"Failed to load SigLIP: {e}")
|
|
return None, None
|
|
|
|
|
|
def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
|
"""Get 768-dim SigLIP embedding for an image."""
|
|
model, processor = get_siglip_model()
|
|
if model is None:
|
|
return None
|
|
|
|
try:
|
|
import torch
|
|
|
|
inputs = processor(images=img_pil, return_tensors="pt")
|
|
device = next(model.parameters()).device
|
|
inputs = {k: v.to(device) for k, v in inputs.items()}
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
return outputs.pooler_output.squeeze().cpu().numpy()
|
|
except Exception as e:
|
|
logger.error(f"Error getting object embedding: {e}")
|
|
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 with Caching
|
|
# =============================================================================
|
|
|
|
|
|
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:
|
|
"""Check if embedding model is available for the given entity type."""
|
|
if entity_type == "face":
|
|
return get_insightface_app() is not None
|
|
model, _ = get_siglip_model()
|
|
return model is not None
|