src/models.py

# src/models.py  —  Enterprise Lens V3
# ════════════════════════════════════════════════════════════════════
# Face Lane  : InsightFace (YuNet detection + ArcFace 512-D encoding)
#              • Replaces DeepFace + RetinaFace + GhostFaceNet entirely
#              • 3-5x faster on CPU, handles small faces (≥20×20 px)
#              • Stores one 512-D vector PER face (not per image)
#              • Each vector carries a base64 face-crop thumbnail
# Object Lane: SigLIP + DINOv2 fused 1536-D (unchanged from V2)
# ════════════════════════════════════════════════════════════════════

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import asyncio
import base64
import functools
import hashlib
import io

import cv2
import numpy as np
import threading
import torch
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")

# ── Constants ─────────────────────────────────────────────────────
YOLO_PERSON_CLASS_ID = 0
MIN_FACE_SIZE        = 20     # minimum face width/height in pixels
MAX_FACES_PER_IMAGE  = 10     # cap faces per image for upload
MAX_CROPS            = 6      # max YOLO object crops per image
MAX_IMAGE_SIZE       = 640    # resize longest edge before inference (V3: 640 vs V2: 512)
FACE_CROP_THUMB_SIZE = 112    # face thumbnail size stored in Pinecone metadata
FACE_CROP_QUALITY    = 75     # JPEG quality for stored thumbnails


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_np: np.ndarray, bbox: list, thumb_size: int = FACE_CROP_THUMB_SIZE) -> str:
    """Crop face from image, resize to thumbnail, return as base64 JPEG string."""
    x, y, w, h = bbox
    x, y = max(0, x), max(0, y)
    # Add 20% padding for more natural face crop
    pad_x = int(w * 0.2)
    pad_y = int(h * 0.2)
    x1 = max(0, x - pad_x)
    y1 = max(0, y - pad_y)
    x2 = min(img_np.shape[1], x + w + pad_x)
    y2 = min(img_np.shape[0], y + h + pad_y)
    face_crop = img_np[y1:y2, x1:x2]
    if face_crop.size == 0:
        return ""
    # Resize to thumbnail
    face_pil  = Image.fromarray(face_crop[..., ::-1])  # BGR → RGB
    face_pil  = face_pil.resize((thumb_size, thumb_size), Image.LANCZOS)
    buf = io.BytesIO()
    face_pil.save(buf, format="JPEG", quality=FACE_CROP_QUALITY)
    return base64.b64encode(buf.getvalue()).decode()


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) ─────────────
        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()

        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 ─────────────────────────
        self.yolo = YOLO("yolo11n-seg.pt")

        # ── Face Lane: InsightFace (YuNet + ArcFace) ─────────────
        self.face_app = None
        print(f"🔍 INSIGHTFACE_AVAILABLE = {INSIGHTFACE_AVAILABLE}")
        if INSIGHTFACE_AVAILABLE:
            try:
                import insightface
                print(f"🔍 InsightFace version: {insightface.__version__}")
                model_name = "buffalo_l" if self.device == "cuda" else "buffalo_sc"
                print(f"🔍 Loading InsightFace model: {model_name}")
                self.face_app = FaceAnalysis(name=model_name)
                self.face_app.prepare(
                    ctx_id=0 if self.device == "cuda" else -1,
                    det_size=(640, 640),
                )
                # Test with a blank image to confirm models loaded
                import numpy as _np
                test_img = _np.zeros((112, 112, 3), dtype=_np.uint8)
                _ = self.face_app.get(test_img)
                print(f"✅ InsightFace ({model_name}) loaded — ArcFace face lane ACTIVE")
            except Exception as e:
                import traceback
                print(f"❌ InsightFace init FAILED: {e}")
                print(traceback.format_exc())
                self.face_app = None
        else:
            print("❌ InsightFace NOT installed — run: pip install insightface onnxruntime")

        self._cache          = {}
        self._cache_maxsize  = 128
        # InsightFace ONNX runtime is NOT thread-safe
        # This lock ensures only one inference runs at a time
        self._face_lock      = threading.Lock()
        print("✅ Models ready!")

    # ── Object Lane batched embedding ────────────────────────────
    def _embed_crops_batch(self, crops: list) -> list:
        if not crops:
            return []
        with torch.no_grad():
            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()

            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))]

    # ── V3 Face detection + encoding ─────────────────────────────
    def _detect_and_encode_faces(self, img_np: np.ndarray) -> list:
        """
        Detect ALL faces in image using InsightFace (YuNet + ArcFace).
        Returns list of dicts:
        {
          "type":       "face",
          "vector":     np.ndarray (512-D ArcFace embedding),
          "face_idx":   int,
          "bbox":       [x, y, w, h],
          "face_crop":  str (base64 JPEG thumbnail),
          "det_score":  float (detection confidence)
        }
        """
        if self.face_app is None:
            print("⚠️  face_app is None — InsightFace not loaded!")
            return []

        try:
            print(f"🔍 Running InsightFace on image shape: {img_np.shape}")
            # InsightFace expects BGR numpy array
            if img_np.shape[2] == 3 and img_np.dtype == np.uint8:
                bgr = img_np[..., ::-1].copy()  # RGB → BGR
            else:
                bgr = img_np.copy()

            with self._face_lock:
                faces = self.face_app.get(bgr)
            print(f"🔍 InsightFace raw detection: {len(faces)} faces found")
            results = []

            for idx, face in enumerate(faces):
                if idx >= MAX_FACES_PER_IMAGE:
                    break

                # Get bounding box
                bbox = face.bbox.astype(int)  # [x1, y1, x2, y2]
                x1, y1, x2, y2 = bbox
                w, h = x2 - x1, y2 - y1

                # Skip tiny faces
                if w < MIN_FACE_SIZE or h < MIN_FACE_SIZE:
                    continue

                # Get ArcFace embedding (already L2-normalised by InsightFace)
                if face.embedding is None:
                    continue
                vec = face.embedding.astype(np.float32)
                # Re-normalise just to be safe
                norm = np.linalg.norm(vec)
                if norm > 0:
                    vec = vec / norm

                # Generate face crop thumbnail for UI
                face_crop_b64 = _crop_to_b64(
                    bgr, [x1, y1, w, h], FACE_CROP_THUMB_SIZE)

                results.append({
                    "type":      "face",
                    "vector":    vec,
                    "face_idx":  idx,
                    "bbox":      [int(x1), int(y1), int(w), int(h)],
                    "face_crop": face_crop_b64,
                    "det_score": float(face.det_score) if hasattr(face, "det_score") else 1.0,
                })

            print(f"👤 Detected {len(results)} face(s) via InsightFace ArcFace")
            return results

        except Exception as e:
            print(f"🟠 InsightFace error: {e} — falling back to object lane")
            return []

    # ── Main process_image ────────────────────────────────────────
    def process_image(
        self,
        image_path: str,
        is_query: bool = False,
        detect_faces: bool = True,
    ) -> list:
        """
        Returns list of vector dicts for upload or search.

        Upload mode (is_query=False):
          - Face vectors include bbox + face_crop for Pinecone metadata
          - Object vectors include full-image + YOLO crops

        Query mode (is_query=True):
          - Same structure — main.py handles grouping for search response
        """
        cache_key = f"{_img_hash(image_path)}_{detect_faces}_{is_query}"
        if cache_key in self._cache:
            print("⚡ Cache hit — skipping inference")
            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 ────────────────────────────────────────────
        if detect_faces:
            # Resize for face detection (640px for small face detection)
            detect_pil = _resize_pil(original_pil, 640)
            detect_np  = np.array(detect_pil)

            face_results = self._detect_and_encode_faces(detect_np)

            if face_results:
                faces_found = True
                # Scale bbox back to original image size if resized
                scale_x = original_pil.width  / detect_pil.width
                scale_y = original_pil.height / detect_pil.height
                for fr in face_results:
                    if scale_x != 1.0 or scale_y != 1.0:
                        bx, by, bw, bh = fr["bbox"]
                        fr["bbox"] = [
                            int(bx * scale_x), int(by * scale_y),
                            int(bw * scale_x), int(bh * scale_y),
                        ]
                    extracted.append(fr)

        # ── OBJECT LANE ──────────────────────────────────────────
        # Always run object lane — even if faces found
        # (image may contain both people and objects)
        crops_pil    = [_resize_pil(original_pil, MAX_IMAGE_SIZE)]  # full-image always
        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())
                    # Skip person crops if face lane already handled them
                    if faces_found and cls_id == YOLO_PERSON_CLASS_ID:
                        print("🔵 PERSON crop skipped — face lane active")
                        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) in one batch …")
        obj_vecs = self._embed_crops_batch(crops)
        for vec in obj_vecs:
            extracted.append({"type": "object", "vector": vec})

        # Cache result
        if len(self._cache) >= self._cache_maxsize:
            oldest = next(iter(self._cache))
            del self._cache[oldest]
        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),
        )