app/Hackathon_setup/face_recognition.py

"""
face_recognition.py
--------------------------
Provides utility functions for:
  - Face detection (via MTCNN)
  - Face similarity computation
  - Face classification using pretrained FaceNet embeddings
"""

import os
import numpy as np
import torch
import torch.nn.functional as F
import joblib
import cv2
from PIL import Image

from .face_recognition_model import FaceRecognitionModel, transform, device, load_model

# =====================================================
# PATH SETUP
# =====================================================
current_path = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(current_path, "face_recognition_model.t7")
CLASSIFIER_PATH = os.path.join(current_path, "team_classifier.joblib")
ENCODER_PATH = os.path.join(current_path, "label_encoder.joblib")

# =====================================================
# LOAD MODELS
# =====================================================
# Load FaceNet model safely with custom loader
model = load_model(MODEL_PATH, pretrained=True)

# Load classifier and label encoder (if available)
clf = joblib.load(CLASSIFIER_PATH) if os.path.exists(CLASSIFIER_PATH) else None
le = joblib.load(ENCODER_PATH) if os.path.exists(ENCODER_PATH) else None

# =====================================================
# FACE DETECTION
# =====================================================
def detected_face(image_bgr):
    """
    Detect a single face in the BGR OpenCV image and return as PIL Image tensor.
    Returns None if no face is detected.
    """
    # Convert BGR → RGB → PIL
    img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    img_pil = Image.fromarray(img_rgb)

    with torch.no_grad():
        face = model.mtcnn(img_pil)
    if face is None:
        return None

    return face

# =====================================================
# SIMILARITY COMPUTATION
# =====================================================
def get_similarity(img1_bgr, img2_bgr):
    """
    Computes cosine similarity between two face images.
    Returns a dictionary with similarity and dissimilarity values.
    """

    emb1 = model.extract_embedding(Image.fromarray(cv2.cvtColor(img1_bgr, cv2.COLOR_BGR2RGB)))
    emb2 = model.extract_embedding(Image.fromarray(cv2.cvtColor(img2_bgr, cv2.COLOR_BGR2RGB)))

    if emb1 is None or emb2 is None:
        return {"similarity": 0.0, "dissimilarity": 1.0}

    emb1_t = torch.tensor(emb1)
    emb2_t = torch.tensor(emb2)

    similarity = F.cosine_similarity(emb1_t, emb2_t).item()
    dissimilarity = round(1 - similarity, 3)

    return {
        "similarity": round(similarity, 3),
        "dissimilarity": dissimilarity
    }

# =====================================================
# CLASS PREDICTION
# =====================================================
def get_face_class(image_bgr):
    """
    Classifies a detected face into one of the known identities.
    Returns: dict(name=<label>, probability=<float>)
    """

    emb = model.extract_embedding(Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)))

    if emb is None or clf is None or le is None:
        return {"name": "Unknown", "probability": 0.0}

    # Predict class using classifier
    pred_idx = clf.predict(emb)[0]
    pred_proba = float(np.max(clf.predict_proba(emb)))

    predicted_class = le.inverse_transform([pred_idx])[0]

    return {"name": predicted_class, "probability": pred_proba}