Spaces:
Sleeping
Sleeping
| # 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), | |
| ) |