src/face_module.py

"""
Face Detection and Recognition Module using InsightFace.

CPU-only inference via ONNX Runtime backend.
- FaceDetection: SCRFD (RetinaFace-based)
- FaceRecognition: ArcFace (512-dim embeddings)
"""

from typing import List, Dict, Optional, Tuple
import numpy as np
import cv2
import os
import sys


class FaceModule:
    """
    Face detection and recognition using InsightFace models.

    Downloads buffalo_l model pack on first use.
    Falls back to ONNX Runtime for CPU inference.
    """

    def __init__(self, det_threshold: float = 0.5, rec_threshold: float = 0.6):
        """
        Initialize face detection and recognition models.

        Args:
            det_threshold: Face detection confidence threshold.
            rec_threshold: Face recognition matching threshold.
        """
        self.det_threshold = det_threshold
        self.rec_threshold = rec_threshold

        # Suppress insightface warnings
        self._model_loaded = False
        self.detector = None
        self.recognizer = None

        try:
            import insightface
            from insightface.model_zoo import get_model
            from insightface.app import FaceAnalysis

            # Initialize FaceAnalysis with buffalo_l (scrfd detection + arcface recognition)
            self.app = FaceAnalysis(
                name='buffalo_l',
                root=os.path.join(os.path.dirname(__file__), '..', 'models', 'insightface'),
                providers=['CPUExecutionProvider']
            )
            self.app.prepare(ctx_id=-1, det_thresh=self.det_threshold)
            self._model_loaded = True
            print(f"[FaceModule] InsightFace loaded successfully (buffalo_l)")
            print(f"[FaceModule] Detection threshold: {self.det_threshold}, Recognition threshold: {self.rec_threshold}")
        except Exception as e:
            print(f"[FaceModule] InsightFace initialization warning: {e}")
            print(f"[FaceModule] Will use lightweight fallback detection")

    def detect_faces(
        self,
        frame: np.ndarray,
        person_bboxes: Optional[List[List[float]]] = None
    ) -> List[Dict]:
        """
        Detect faces in a frame, optionally within person bounding boxes.

        Args:
            frame: Input image (H, W, C) in BGR format.
            person_bboxes: Optional list of person bboxes [x1,y1,x2,y2] to constrain search.

        Returns:
            List of face dicts with keys: face_bbox, landmarks, embedding, det_score.
        """
        if not self._model_loaded or frame is None or frame.size == 0:
            return []

        try:
            all_faces = self.app.get(frame)

            if person_bboxes:
                # Filter faces to those within person bounding boxes
                filtered_faces = []
                for face in all_faces:
                    face_bbox = face.bbox.astype(float).tolist()
                    face_cx = (face_bbox[0] + face_bbox[2]) / 2
                    face_cy = (face_bbox[1] + face_bbox[3]) / 2

                    for pbox in person_bboxes:
                        if pbox[0] <= face_cx <= pbox[2] and pbox[1] <= face_cy <= pbox[3]:
                            filtered_faces.append(face)
                            break
                faces = filtered_faces
            else:
                faces = all_faces

            results = []
            for face in faces:
                result = {
                    'face_bbox': face.bbox.astype(float).tolist(),
                    'landmarks': face.landmark.astype(float).tolist() if face.landmark is not None else None,
                    'embedding': face.normed_embedding.astype(np.float32) if face.normed_embedding is not None else None,
                    'det_score': float(face.det_score),
                }
                results.append(result)

            return results

        except Exception as e:
            print(f"[FaceModule] Detection error: {e}")
            return []

    def get_embedding(self, face_crop: np.ndarray) -> Optional[np.ndarray]:
        """
        Get face embedding from a cropped face image.

        Args:
            face_crop: Cropped face image (BGR).

        Returns:
            512-dim embedding vector or None if no face found.
        """
        if not self._model_loaded or face_crop is None or face_crop.size == 0:
            return None

        try:
            faces = self.app.get(face_crop)
            if len(faces) > 0:
                return faces[0].normed_embedding.astype(np.float32)
            return None
        except Exception as e:
            print(f"[FaceModule] Embedding error: {e}")
            return None

    def compare_embeddings(self, emb1: np.ndarray, emb2: np.ndarray) -> float:
        """
        Compare two face embeddings using cosine similarity.

        Args:
            emb1: First embedding vector.
            emb2: Second embedding vector.

        Returns:
            Cosine similarity score (0-1).
        """
        if emb1 is None or emb2 is None:
            return 0.0

        emb1 = emb1.flatten()
        emb2 = emb2.flatten()

        norm1 = np.linalg.norm(emb1)
        norm2 = np.linalg.norm(emb2)

        if norm1 < 1e-10 or norm2 < 1e-10:
            return 0.0

        similarity = float(np.dot(emb1, emb2) / (norm1 * norm2))
        return max(0.0, min(1.0, similarity))

    def is_loaded(self) -> bool:
        """Check if models are loaded."""
        return self._model_loaded


if __name__ == "__main__":
    # Quick test
    import numpy as np
    fm = FaceModule()
    test_frame = np.zeros((480, 640, 3), dtype=np.uint8)
    faces = fm.detect_faces(test_frame)
    print(f"Face detection test: {len(faces)} faces found")
    print("FaceModule OK!")