encoders.py

import os
import pickle
from typing import Dict, Any, List

import face_recognition
import numpy as np

from logger import get_logger

logger = get_logger(__name__)


def encode_known_faces(known_dir: str, output_path: str = os.path.join("data", "encodings.pickle"),
                       detection_model: str = "hog", force: bool = False) -> Dict[str, Any]:
    """
    Encode all face images in known_dir and save to output_path.
    Returns dict with 'encodings' and 'names' lists.
    """
    if os.path.exists(output_path) and not force:
        logger.info(f"Encodings already exist at {output_path}. Load them instead of recomputing.")
        with open(output_path, "rb") as f:
            return pickle.load(f)

    encodings: List[np.ndarray] = []
    names: List[str] = []
    files: List[str] = []
    rollnos: List[str] = []

    if not os.path.isdir(known_dir):
        raise FileNotFoundError(f"Known directory not found: {known_dir}")

    for fname in sorted(os.listdir(known_dir)):
        if not fname.lower().endswith((".jpg", ".jpeg", ".png")):
            continue
        path = os.path.join(known_dir, fname)
        logger.info(f"Processing {path}")
        image = face_recognition.load_image_file(path)
        face_locations = face_recognition.face_locations(image, model=detection_model)
        if len(face_locations) == 0:
            logger.warning(f"No face detected in {fname}; skipping.")
            continue
        # If multiple faces in a single file, we take first encoding (you can change this later)
        face_encs = face_recognition.face_encodings(image, face_locations)
        if len(face_encs) == 0:
            logger.warning(f"Could not compute encodings for {fname}; skipping.")
            continue

        # name extraction: use <id>_<name>.jpg or name.jpg
        name = os.path.splitext(fname)[0]
        roll_no = None
        # Expect filename format: rollno_name.jpg (e.g., 12_Ramesh.jpg)
        if "_" in name:
            roll_no, name = name.split("_", 1)
        else:
            name = name
            roll_no = "UNKNOWN"  # fallback if no rollno in filename

        encodings.append(face_encs[0])
        names.append(name)
        files.append(fname)
        rollnos.append(roll_no)

    data = {"encodings": encodings, "names": names, "files": files, "rollnos": rollnos}
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    with open(output_path, "wb") as f:
        pickle.dump(data, f)
    logger.info(f"Saved {len(names)} encodings to {output_path}")
    return data

if __name__ == "__main__":
    encode_known_faces(known_dir = os.path.join("data", "known_faces"), force = True)