193 lines
6.6 KiB
Python
193 lines
6.6 KiB
Python
"""
|
|
Unified embedding interface for faces and objects.
|
|
|
|
- Faces: InsightFace (ArcFace/Buffalo_L)
|
|
- Objects: SigLIP (Vision Transformer via transformers)
|
|
"""
|
|
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
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)
|
|
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
|
|
|
|
|
|
# =============================================================================
|
|
# Unified Interface
|
|
# =============================================================================
|
|
|
|
|
|
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 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
|