src/face_processor.py

"""
Face Processor - Detección y alineación de rostros
"""

import cv2
import numpy as np
from mtcnn import MTCNN
from PIL import Image
from loguru import logger


class FaceProcessor:
    """
    Procesa imágenes para detectar, alinear y normalizar rostros.
    """
    
    def __init__(self):
        """Inicializa el detector MTCNN"""
        logger.info("Inicializando MTCNN...")
        self.detector = MTCNN()
        logger.success("MTCNN inicializado")
    
    def align_face(self, image):
        """
        Detecta y alinea el rostro en la imagen.
        
        Args:
            image: Imagen PIL o numpy array (RGB)
            
        Returns:
            Rostro alineado y normalizado (160x160) o None si no se detecta
        """
        # Convertir PIL a numpy si es necesario
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Asegurar que está en RGB
        if len(image.shape) == 2:  # Grayscale
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        elif image.shape[2] == 4:  # RGBA
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        
        # Detectar rostros
        faces = self.detector.detect_faces(image)
        
        if len(faces) == 0:
            logger.warning("No se detectó ningún rostro")
            return None
        
        # Tomar el rostro más grande (más probable que sea el principal)
        face = max(faces, key=lambda x: x['box'][2] * x['box'][3])
        
        # Extraer keypoints
        keypoints = face['keypoints']
        left_eye = keypoints['left_eye']
        right_eye = keypoints['right_eye']
        
        # Calcular ángulo de rotación para alinear horizontalmente
        dY = right_eye[1] - left_eye[1]
        dX = right_eye[0] - left_eye[0]
        angle = np.degrees(np.arctan2(dY, dX))
        
        # Rotar imagen
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        aligned = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC)
        
        # Recortar rostro con margen
        x, y, width, height = face['box']
        margin = int(min(width, height) * 0.2)  # 20% de margen
        
        x1 = max(0, x - margin)
        y1 = max(0, y - margin)
        x2 = min(w, x + width + margin)
        y2 = min(h, y + height + margin)
        
        face_crop = aligned[y1:y2, x1:x2]
        
        # Resize a 160x160 (estándar FaceNet)
        try:
            face_resized = cv2.resize(face_crop, (160, 160), interpolation=cv2.INTER_AREA)
            logger.debug(f"Rostro detectado y alineado: {face_resized.shape}")
            return face_resized
        except Exception as e:
            logger.error(f"Error al resize: {e}")
            return None