visual-search-api2 / src /models.py
AdarshDRC's picture
fix : search engine
1feca1e
Raw
History Blame
29 kB
# src/models.py — Enterprise Lens V4
# ════════════════════════════════════════════════════════════════════
# Face Lane : InsightFace SCRFD-10GF + ArcFace-R100 (buffalo_l)
# + AdaFace IR-50 (WebFace4M) fused → 1024-D vector
# • det_size=(1280,1280) — catches small/group faces
# • Quality gate: det_score ≥ 0.60, face_px ≥ 40
# • Multi-scale: runs detection at 2 scales, merges
# • Stores one 1024-D vector PER face
# • Each vector carries base64 face-crop thumbnail
# • face_quality_score + face_width_px in metadata
#
# Object Lane: SigLIP + DINOv2 fused 1536-D (unchanged from V3)
# ════════════════════════════════════════════════════════════════════
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import asyncio
import base64
import functools
import hashlib
import io
import threading
import traceback
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel, AutoProcessor
from ultralytics import YOLO
# ── InsightFace ───────────────────────────────────────────────────
try:
import insightface
from insightface.app import FaceAnalysis
INSIGHTFACE_AVAILABLE = True
except ImportError:
INSIGHTFACE_AVAILABLE = False
print("⚠️ insightface not installed — face lane disabled")
print(" Run: pip install insightface onnxruntime-silicon (mac)")
print(" pip install insightface onnxruntime (linux/win)")
# ── AdaFace ──────────────────────────────────────────────────────
# AdaFace IR-50 MS1MV2 (CVPR 2022) — quality-adaptive margin loss
# Repo : minchul/cvlface_adaface_ir50_ms1mv2 (HuggingFace)
# Loaded : AutoModel + trust_remote_code=True (custom_code repo)
# Needs : HF_TOKEN env var set in HF Space secrets
try:
import shutil as _shutil
from huggingface_hub import hf_hub_download
from transformers import AutoModel as _HF_AutoModel
ADAFACE_WEIGHTS_AVAILABLE = True
except ImportError:
ADAFACE_WEIGHTS_AVAILABLE = False
print("⚠️ huggingface_hub / transformers not installed — AdaFace fusion disabled")
# ── Constants ─────────────────────────────────────────────────────
YOLO_PERSON_CLASS_ID = 0
MIN_FACE_SIZE = 40 # V4: stricter — tiny faces embed poorly
MAX_FACES_PER_IMAGE = 12 # slightly higher cap for group photos
MAX_CROPS = 6 # max YOLO object crops per image
MAX_IMAGE_SIZE = 640 # object lane longest edge
DET_SIZE_PRIMARY = (1280, 1280) # V4: 1280 for small-face detection
DET_SIZE_SECONDARY = (640, 640) # fallback / 2nd scale
FACE_CROP_THUMB_SIZE = 112 # face thumbnail for Pinecone metadata
FACE_CROP_QUALITY = 80 # JPEG quality for thumbnails
FACE_QUALITY_GATE = 0.60 # minimum det_score to accept a face
FACE_DIM = 512 # ArcFace embedding dimension
ADAFACE_DIM = 512 # AdaFace embedding dimension
FUSED_FACE_DIM = 1024 # ArcFace + AdaFace concatenated
# ════════════════════════════════════════════════════════════════
# Utility functions
# ════════════════════════════════════════════════════════════════
def _resize_pil(img: Image.Image, max_side: int = MAX_IMAGE_SIZE) -> Image.Image:
w, h = img.size
if max(w, h) <= max_side:
return img
scale = max_side / max(w, h)
return img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
def _img_hash(image_path: str) -> str:
h = hashlib.md5()
with open(image_path, "rb") as f:
h.update(f.read(65536))
return h.hexdigest()
def _crop_to_b64(
img_bgr: np.ndarray,
x1: int, y1: int, x2: int, y2: int,
thumb_size: int = FACE_CROP_THUMB_SIZE,
) -> str:
"""Crop face from BGR image with 20% padding, return base64 JPEG thumbnail."""
H, W = img_bgr.shape[:2]
w, h = x2 - x1, y2 - y1
pad_x = int(w * 0.20)
pad_y = int(h * 0.20)
cx1 = max(0, x1 - pad_x)
cy1 = max(0, y1 - pad_y)
cx2 = min(W, x2 + pad_x)
cy2 = min(H, y2 + pad_y)
crop = img_bgr[cy1:cy2, cx1:cx2]
if crop.size == 0:
return ""
pil = Image.fromarray(crop[:, :, ::-1]) # BGR → RGB
pil = pil.resize((thumb_size, thumb_size), Image.LANCZOS)
buf = io.BytesIO()
pil.save(buf, format="JPEG", quality=FACE_CROP_QUALITY)
return base64.b64encode(buf.getvalue()).decode()
def _face_crop_for_adaface(
img_bgr: np.ndarray,
x1: int, y1: int, x2: int, y2: int,
) -> np.ndarray:
"""
Crop and normalise face for AdaFace IR-50 input.
Returns float32 numpy array (3, 112, 112) normalised to [-1, 1].
"""
H, W = img_bgr.shape[:2]
w, h = x2 - x1, y2 - y1
pad_x = int(w * 0.10)
pad_y = int(h * 0.10)
cx1 = max(0, x1 - pad_x)
cy1 = max(0, y1 - pad_y)
cx2 = min(W, x2 + pad_x)
cy2 = min(H, y2 + pad_y)
crop = img_bgr[cy1:cy2, cx1:cx2]
if crop.size == 0:
return None
rgb = crop[:, :, ::-1].copy() # BGR → RGB
pil = Image.fromarray(rgb).resize((112, 112), Image.LANCZOS)
arr = np.array(pil, dtype=np.float32) / 255.0
arr = (arr - 0.5) / 0.5 # normalise [-1, 1]
return arr.transpose(2, 0, 1) # HWC → CHW
# ════════════════════════════════════════════════════════════════
# AIModelManager — V4
# ════════════════════════════════════════════════════════════════
class AIModelManager:
def __init__(self):
self.device = (
"cuda" if torch.cuda.is_available()
else ("mps" if torch.backends.mps.is_available() else "cpu")
)
print(f"🚀 Loading models onto: {self.device.upper()}...")
# ── Object Lane: SigLIP + DINOv2 (unchanged) ─────────────
print("📦 Loading SigLIP...")
self.siglip_processor = AutoProcessor.from_pretrained(
"google/siglip-base-patch16-224", use_fast=True)
self.siglip_model = AutoModel.from_pretrained(
"google/siglip-base-patch16-224").to(self.device).eval()
print("📦 Loading DINOv2...")
self.dinov2_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
self.dinov2_model = AutoModel.from_pretrained(
"facebook/dinov2-base").to(self.device).eval()
if self.device == "cuda":
self.siglip_model = self.siglip_model.half()
self.dinov2_model = self.dinov2_model.half()
# ── YOLO for object segmentation ─────────────────────────
print("📦 Loading YOLO11n-seg...")
self.yolo = YOLO("yolo11n-seg.pt")
# ── Face Lane: InsightFace SCRFD + ArcFace-R100 ───────────
# V4: ALWAYS use buffalo_l (SCRFD-10GF + ArcFace-R100)
# even on CPU — accuracy matters more than speed here.
# det_size=1280 catches faces as small as ~10px in source.
self.face_app = None
if INSIGHTFACE_AVAILABLE:
try:
print("📦 Loading InsightFace buffalo_l (SCRFD-10GF + ArcFace-R100)...")
self.face_app = FaceAnalysis(
name="buffalo_l",
providers=(
["CUDAExecutionProvider", "CPUExecutionProvider"]
if self.device == "cuda"
else ["CPUExecutionProvider"]
),
)
self.face_app.prepare(
ctx_id=0 if self.device == "cuda" else -1,
det_size=DET_SIZE_PRIMARY, # 1280×1280 — key for small faces
)
# Warmup
test_img = np.zeros((112, 112, 3), dtype=np.uint8)
self.face_app.get(test_img)
print("✅ InsightFace buffalo_l loaded — SCRFD+ArcFace face lane ACTIVE")
print(f" det_size={DET_SIZE_PRIMARY} | quality_gate={FACE_QUALITY_GATE}")
except Exception as e:
print(f"❌ InsightFace init FAILED: {e}")
print(traceback.format_exc())
self.face_app = None
else:
print("❌ InsightFace NOT installed")
# ── AdaFace IR-50 (CVPR 2022) — quality-adaptive fusion ───
# Fused with ArcFace → 1024-D face vector
# Weights: adaface_ir50_webface4m.ckpt from HuggingFace
self.adaface_model = None
self._load_adaface()
# Thread safety for ONNX
self._face_lock = threading.Lock()
self._cache = {}
self._cache_maxsize = 128
adaface_status = "FULL FUSION u2705" if self.adaface_model else "ZERO-PADDED u26a0ufe0f (AdaFace weights missing)"
print("")
print("u2705 Enterprise Lens V4 u2014 Models Ready")
print(f" Device : {self.device.upper()}")
print(f" InsightFace : buffalo_l (SCRFD-10GF + ArcFace-R100)")
print(f" AdaFace : {adaface_status}")
print(f" Face vector dim : {FUSED_FACE_DIM} <- enterprise-faces MUST be {FUSED_FACE_DIM}-D")
print(f" Object vector dim : 1536 <- enterprise-objects MUST be 1536-D")
print(f" Quality gate : det_score >= {FACE_QUALITY_GATE}, face_px >= {MIN_FACE_SIZE}")
print(f" Detection size : {DET_SIZE_PRIMARY}")
print("")
def _load_adaface(self):
"""
Load AdaFace IR-50 MS1MV2 from HuggingFace.
Repo : minchul/cvlface_adaface_ir50_ms1mv2
Method : AutoModel + trust_remote_code (repo has custom_code)
Token : HF_TOKEN env var (required for custom_code repos)
Output : 512-D L2-normalised embedding per face crop
"""
if not ADAFACE_WEIGHTS_AVAILABLE:
print("⚠️ AdaFace skipped — huggingface_hub / transformers not installed")
return
import os, sys
REPO_ID = "minchul/cvlface_adaface_ir50_ms1mv2"
HF_TOKEN = os.getenv("HF_TOKEN", None)
CACHE_PATH = os.path.expanduser("~/.cvlface_cache/minchul/cvlface_adaface_ir50_ms1mv2")
try:
print("📦 Loading AdaFace IR-50 MS1MV2 from HuggingFace...")
if HF_TOKEN:
print(" HF_TOKEN found ✅")
else:
print(" ⚠️ HF_TOKEN not set — may fail on gated/custom_code repos")
# ── Step 1: Download all repo files ──────────────────
os.makedirs(CACHE_PATH, exist_ok=True)
# Download files.txt manifest first
files_txt = os.path.join(CACHE_PATH, "files.txt")
if not os.path.exists(files_txt):
hf_hub_download(
repo_id=REPO_ID, filename="files.txt",
token=HF_TOKEN, local_dir=CACHE_PATH,
local_dir_use_symlinks=False,
)
# Read manifest and download each listed file
with open(files_txt, "r") as f:
extra_files = [x.strip() for x in f.read().split("\n") if x.strip()]
for fname in extra_files + ["config.json", "wrapper.py", "model.safetensors"]:
fpath = os.path.join(CACHE_PATH, fname)
if not os.path.exists(fpath):
print(f" Downloading {fname}...")
hf_hub_download(
repo_id=REPO_ID, filename=fname,
token=HF_TOKEN, local_dir=CACHE_PATH,
local_dir_use_symlinks=False,
)
# ── Step 2: Load model from local cache ──────────────
# Must chdir + add to sys.path because the repo uses
# trust_remote_code with relative imports in wrapper.py
cwd = os.getcwd()
os.chdir(CACHE_PATH)
sys.path.insert(0, CACHE_PATH)
try:
model = _HF_AutoModel.from_pretrained(
CACHE_PATH,
trust_remote_code=True,
token=HF_TOKEN,
)
finally:
os.chdir(cwd)
if CACHE_PATH in sys.path:
sys.path.remove(CACHE_PATH)
model = model.to(self.device).eval()
if self.device == "cuda":
model = model.half()
# ── Step 3: Verify output shape ───────────────────────
with torch.no_grad():
dummy = torch.zeros(1, 3, 112, 112).to(self.device)
out = model(dummy)
# Model may return tensor directly or an object with .embedding
out_vec = out if isinstance(out, torch.Tensor) else out.embedding
out_dim = out_vec.shape[-1]
if out_dim != ADAFACE_DIM:
raise ValueError(
f"AdaFace output dim={out_dim}, expected {ADAFACE_DIM}")
self.adaface_model = model
print(f"✅ AdaFace IR-50 MS1MV2 loaded — output dim={out_dim} — 1024-D fusion ACTIVE")
except Exception as e:
print(f"⚠️ AdaFace load failed: {e}")
print(f" Detail: {traceback.format_exc()[-500:]}")
print(" Falling back to ArcFace-only (zero-padded to 1024-D)")
self.adaface_model = None
# ── Object Lane: batched SigLIP + DINOv2 embedding ───────────
def _embed_crops_batch(self, crops: list) -> list:
"""Embed a list of PIL images → list of 1536-D numpy arrays."""
if not crops:
return []
with torch.no_grad():
# SigLIP
sig_in = self.siglip_processor(images=crops, return_tensors="pt", padding=True)
sig_in = {k: v.to(self.device) for k, v in sig_in.items()}
if self.device == "cuda":
sig_in = {k: v.half() if v.dtype == torch.float32 else v
for k, v in sig_in.items()}
sig_out = self.siglip_model.get_image_features(**sig_in)
if hasattr(sig_out, "image_embeds"): sig_out = sig_out.image_embeds
elif isinstance(sig_out, tuple): sig_out = sig_out[0]
sig_vecs = F.normalize(sig_out.float(), p=2, dim=1).cpu()
# DINOv2
dino_in = self.dinov2_processor(images=crops, return_tensors="pt")
dino_in = {k: v.to(self.device) for k, v in dino_in.items()}
if self.device == "cuda":
dino_in = {k: v.half() if v.dtype == torch.float32 else v
for k, v in dino_in.items()}
dino_out = self.dinov2_model(**dino_in)
dino_vecs = F.normalize(
dino_out.last_hidden_state[:, 0, :].float(), p=2, dim=1).cpu()
fused = F.normalize(torch.cat([sig_vecs, dino_vecs], dim=1), p=2, dim=1)
return [fused[i].numpy() for i in range(len(crops))]
# ── AdaFace embedding for a single face crop ─────────────────
def _adaface_embed(self, face_arr_chw: np.ndarray) -> np.ndarray:
"""
Run AdaFace IR-50 MS1MV2 on a preprocessed (3,112,112) float32 array.
Input : CHW float32, normalised to [-1, 1]
Output: 512-D L2-normalised numpy embedding, or None on failure.
The cvlface model may return a tensor directly or an object
with an .embedding attribute — both cases handled.
"""
if self.adaface_model is None or face_arr_chw is None:
return None
try:
t = torch.from_numpy(face_arr_chw).unsqueeze(0) # (1,3,112,112)
t = t.to(self.device)
if self.device == "cuda":
t = t.half()
with torch.no_grad():
out = self.adaface_model(t)
# Handle both raw tensor and object-with-embedding outputs
emb = out if isinstance(out, torch.Tensor) else out.embedding
emb = F.normalize(emb.float(), p=2, dim=1)
return emb[0].cpu().numpy()
except Exception as e:
print(f"⚠️ AdaFace inference error: {e}")
return None
# ── V4 Face detection + dual encoding ────────────────────────
def _detect_and_encode_faces(self, img_np: np.ndarray) -> list:
"""
Detect ALL faces using InsightFace SCRFD-10GF at 1280px.
For each face:
- ArcFace-R100 embedding (512-D, from InsightFace)
- AdaFace IR-50 embedding (512-D, fused quality-adaptive)
- Concatenate + L2-normalise → 1024-D final vector
- Quality gate: det_score ≥ 0.60, face width ≥ 40px
- Base64 thumbnail stored for UI
Returns list of dicts with keys:
type, vector (1024-D or 512-D), face_idx, bbox,
face_crop, det_score, face_quality, face_width_px
"""
if self.face_app is None:
print("⚠️ face_app is None — InsightFace not loaded")
return []
try:
# InsightFace expects BGR
if img_np.dtype != np.uint8:
img_np = (img_np * 255).astype(np.uint8)
bgr = img_np[:, :, ::-1].copy() if img_np.shape[2] == 3 else img_np.copy()
print(f"🔍 SCRFD detection on {bgr.shape[1]}×{bgr.shape[0]} image...")
with self._face_lock:
faces = self.face_app.get(bgr)
print(f" Raw detections: {len(faces)}")
results = []
accepted = 0
for idx, face in enumerate(faces):
if accepted >= MAX_FACES_PER_IMAGE:
break
# ── Bounding box ──────────────────────────────────
bbox_raw = face.bbox.astype(int)
x1, y1, x2, y2 = bbox_raw
x1 = max(0, x1); y1 = max(0, y1)
x2 = min(bgr.shape[1], x2); y2 = min(bgr.shape[0], y2)
w, h = x2 - x1, y2 - y1
if w <= 0 or h <= 0:
continue
# ── Quality gate 1: minimum size ──────────────────
if w < MIN_FACE_SIZE or h < MIN_FACE_SIZE:
print(f" Face {idx}: SKIP — too small ({w}×{h}px)")
continue
# ── Quality gate 2: detection confidence ──────────
det_score = float(face.det_score) if hasattr(face, "det_score") else 1.0
if det_score < FACE_QUALITY_GATE:
print(f" Face {idx}: SKIP — low det_score ({det_score:.3f})")
continue
# ── ArcFace embedding (from InsightFace) ──────────
if face.embedding is None:
continue
arcface_vec = face.embedding.astype(np.float32)
n = np.linalg.norm(arcface_vec)
if n > 0:
arcface_vec = arcface_vec / n
# ── AdaFace embedding (quality-adaptive) ──────────
face_chw = _face_crop_for_adaface(bgr, x1, y1, x2, y2)
adaface_vec = self._adaface_embed(face_chw)
# ── Fuse: ArcFace + AdaFace → 1024-D ─────────────
# ALWAYS output FUSED_FACE_DIM (1024) so Pinecone index
# dimension never mismatches, regardless of AdaFace status.
if adaface_vec is not None:
# Full fusion: ArcFace(512) + AdaFace(512) → 1024-D
fused_raw = np.concatenate([arcface_vec, adaface_vec])
else:
# AdaFace unavailable — pad with zeros to maintain 1024-D
# The ArcFace half still carries full identity signal;
# zero padding is neutral and doesn't corrupt similarity.
print(" ⚠️ AdaFace unavailable — padding to 1024-D")
fused_raw = np.concatenate([arcface_vec,
np.zeros(ADAFACE_DIM, dtype=np.float32)])
n2 = np.linalg.norm(fused_raw)
final_vec = (fused_raw / n2) if n2 > 0 else fused_raw
vec_dim = FUSED_FACE_DIM # always 1024
# ── Face crop thumbnail for UI ─────────────────────
face_crop_b64 = _crop_to_b64(bgr, x1, y1, x2, y2)
results.append({
"type": "face",
"vector": final_vec,
"vec_dim": vec_dim,
"face_idx": accepted,
"bbox": [int(x1), int(y1), int(w), int(h)],
"face_crop": face_crop_b64,
"det_score": det_score,
"face_quality": det_score, # alias for metadata
"face_width_px": int(w),
})
accepted += 1
print(f" Face {idx}: ACCEPTED — {w}×{h}px | "
f"det={det_score:.3f} | dim={vec_dim}")
print(f"👤 {accepted} face(s) passed quality gate")
return results
except Exception as e:
print(f"🟠 InsightFace error: {e}")
print(traceback.format_exc()[-600:])
return []
# ── Main process_image ────────────────────────────────────────
def process_image(
self,
image_path: str,
is_query: bool = False,
detect_faces: bool = True,
) -> list:
"""
Full pipeline for one image.
Returns list of vector dicts:
Face: {type, vector (1024-D), face_idx, bbox, face_crop,
det_score, face_quality, face_width_px}
Object: {type, vector (1536-D)}
V4 changes vs V3:
- SCRFD at 1280px (not 640) — catches small/group faces
- buffalo_l always (not buffalo_sc on CPU)
- ArcFace + AdaFace fused 1024-D vectors
- Quality gate: det_score ≥ 0.60, width ≥ 40px
- Multi-scale: detect at 1280, retry at 640 if 0 faces found
"""
cache_key = f"{_img_hash(image_path)}_{detect_faces}_{is_query}"
if cache_key in self._cache:
print("⚡ Cache hit")
return self._cache[cache_key]
extracted = []
original_pil = Image.open(image_path).convert("RGB")
img_np = np.array(original_pil) # RGB uint8
faces_found = False
# ════════════════════════════════════════════════════════
# FACE LANE
# V4: Run at full resolution (up to 1280px) to catch small
# faces in group photos. If 0 faces detected, retry at
# the original resolution (multi-scale fallback).
# ════════════════════════════════════════════════════════
if detect_faces and self.face_app is not None:
# Scale 1: resize longest edge to 1280 for detection
detect_pil_1280 = _resize_pil(original_pil, 1280)
detect_np_1280 = np.array(detect_pil_1280)
face_results = self._detect_and_encode_faces(detect_np_1280)
# Scale 2: if nothing found, try original resolution
# (sometimes resizing DOWN helps when image is already small)
if not face_results and max(original_pil.size) < 1280:
print("🔄 Multi-scale fallback: retrying at original resolution")
face_results = self._detect_and_encode_faces(img_np)
if face_results:
faces_found = True
# Scale bboxes back to original-image coordinates
sx = original_pil.width / detect_pil_1280.width
sy = original_pil.height / detect_pil_1280.height
for fr in face_results:
if sx != 1.0 or sy != 1.0:
bx, by, bw, bh = fr["bbox"]
fr["bbox"] = [
int(bx * sx), int(by * sy),
int(bw * sx), int(bh * sy),
]
extracted.append(fr)
# ════════════════════════════════════════════════════════
# OBJECT LANE
# Always runs — even when faces are found.
# PERSON-class YOLO crops are skipped when faces active
# to avoid double-counting people.
# ════════════════════════════════════════════════════════
crops_pil = [_resize_pil(original_pil, MAX_IMAGE_SIZE)] # full image
yolo_results = self.yolo(image_path, conf=0.5, verbose=False)
for r in yolo_results:
if r.masks is not None:
for seg_idx, mask_xy in enumerate(r.masks.xy):
cls_id = int(r.boxes.cls[seg_idx].item())
if faces_found and cls_id == YOLO_PERSON_CLASS_ID:
continue
polygon = np.array(mask_xy, dtype=np.int32)
if len(polygon) < 3:
continue
x, y, w, h = cv2.boundingRect(polygon)
if w < 30 or h < 30:
continue
crop = original_pil.crop((x, y, x + w, y + h))
crops_pil.append(crop)
if len(crops_pil) >= MAX_CROPS + 1:
break
elif r.boxes is not None:
for box in r.boxes:
cls_id = int(box.cls.item())
if faces_found and cls_id == YOLO_PERSON_CLASS_ID:
continue
x1, y1, x2, y2 = box.xyxy[0].tolist()
if (x2 - x1) < 30 or (y2 - y1) < 30:
continue
crop = original_pil.crop((x1, y1, x2, y2))
crops_pil.append(crop)
if len(crops_pil) >= MAX_CROPS + 1:
break
crops = [_resize_pil(c, MAX_IMAGE_SIZE) for c in crops_pil]
print(f"🧠 Embedding {len(crops)} object crop(s)...")
obj_vecs = self._embed_crops_batch(crops)
for vec in obj_vecs:
extracted.append({"type": "object", "vector": vec})
# Cache
if len(self._cache) >= self._cache_maxsize:
del self._cache[next(iter(self._cache))]
self._cache[cache_key] = extracted
return extracted
async def process_image_async(
self,
image_path: str,
is_query: bool = False,
detect_faces: bool = True,
) -> list:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
functools.partial(self.process_image, image_path, is_query, detect_faces),
)