face_detector.py

"""
Face Detection Module
Uses OpenCV's DNN face detector for robust face detection
"""

import cv2
import numpy as np
import logging
from typing import List, Tuple, Optional

logger = logging.getLogger(__name__)


class FaceDetector:
    """
    Face detector using OpenCV DNN module with pre-trained models
    """
    
    def __init__(self, confidence_threshold: float = 0.5):
        """
        Initialize face detector
        
        Args:
            confidence_threshold: Minimum confidence for face detection
        """
        self.confidence_threshold = confidence_threshold
        self.net = None
        self._load_model()
    
    def _load_model(self):
        """Load pre-trained face detection model"""
        try:
            # Using OpenCV's DNN face detector (Caffe model)
            # This is a lightweight and efficient model
            prototxt_path = "models/deploy.prototxt"
            model_path = "models/res10_300x300_ssd_iter_140000.caffemodel"
            
            # Try to load from local files first
            try:
                self.net = cv2.dnn.readNetFromCaffe(prototxt_path, model_path)
                logger.info("✓ Loaded face detection model from local files")
            except:
                # Fallback: use Haar Cascade (built-in to OpenCV)
                logger.warning("DNN model not found, using Haar Cascade fallback")
                self.face_cascade = cv2.CascadeClassifier(
                    cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
                )
                self.net = None
        
        except Exception as e:
            logger.error(f"Error loading face detection model: {e}")
            # Use Haar Cascade as fallback
            self.face_cascade = cv2.CascadeClassifier(
                cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
            )
            self.net = None
    
    def detect_faces(self, image: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """
        Detect faces in an image
        
        Args:
            image: Input image (BGR format)
        
        Returns:
            List of face bounding boxes [(x, y, w, h), ...]
        """
        if image is None or image.size == 0:
            logger.warning("Empty image provided")
            return []
        
        try:
            if self.net is not None:
                return self._detect_dnn(image)
            else:
                return self._detect_haar(image)
        except Exception as e:
            logger.error(f"Face detection error: {e}")
            return []
    
    def _detect_dnn(self, image: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """Detect faces using DNN model"""
        h, w = image.shape[:2]
        
        # Prepare blob for DNN
        blob = cv2.dnn.blobFromImage(
            cv2.resize(image, (300, 300)),
            1.0,
            (300, 300),
            (104.0, 177.0, 123.0)
        )
        
        self.net.setInput(blob)
        detections = self.net.forward()
        
        faces = []
        for i in range(detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            
            if confidence > self.confidence_threshold:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (x1, y1, x2, y2) = box.astype("int")
                
                # Convert to (x, y, w, h) format
                x = max(0, x1)
                y = max(0, y1)
                w = min(image.shape[1] - x, x2 - x1)
                h = min(image.shape[0] - y, y2 - y1)
                
                if w > 0 and h > 0:
                    faces.append((x, y, w, h))
        
        return faces
    
    def _detect_haar(self, image: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """Detect faces using Haar Cascade"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        faces = self.face_cascade.detectMultiScale(
            gray,
            scaleFactor=1.1,
            minNeighbors=5,
            minSize=(30, 30),
            flags=cv2.CASCADE_SCALE_IMAGE
        )
        
        return [tuple(face) for face in faces]
    
    def assess_face_quality(self, image: np.ndarray, face: Tuple[int, int, int, int]) -> float:
        """
        Assess the quality of detected face
        
        Args:
            image: Input image
            face: Face bounding box (x, y, w, h)
        
        Returns:
            Quality score between 0 and 1
        """
        try:
            x, y, w, h = face
            face_roi = image[y:y+h, x:x+w]
            
            if face_roi.size == 0:
                return 0.0
            
            # Convert to grayscale
            gray_face = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
            
            # 1. Size score (larger faces are better)
            size_score = min(1.0, (w * h) / (image.shape[0] * image.shape[1] * 0.5))
            
            # 2. Sharpness score (using Laplacian variance)
            laplacian_var = cv2.Laplacian(gray_face, cv2.CV_64F).var()
            sharpness_score = min(1.0, laplacian_var / 500.0)
            
            # 3. Brightness score
            mean_brightness = np.mean(gray_face)
            brightness_score = 1.0 - abs(mean_brightness - 127.5) / 127.5
            
            # 4. Contrast score
            contrast = gray_face.std()
            contrast_score = min(1.0, contrast / 64.0)
            
            # Weighted average
            quality = (
                size_score * 0.3 +
                sharpness_score * 0.3 +
                brightness_score * 0.2 +
                contrast_score * 0.2
            )
            
            return quality
            
        except Exception as e:
            logger.error(f"Quality assessment error: {e}")
            return 0.0
    
    def draw_faces(self, image: np.ndarray, faces: List[Tuple[int, int, int, int]]) -> np.ndarray:
        """
        Draw bounding boxes around detected faces
        
        Args:
            image: Input image
            faces: List of face bounding boxes
        
        Returns:
            Image with drawn bounding boxes
        """
        output = image.copy()
        
        for (x, y, w, h) in faces:
            cv2.rectangle(output, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(
                output,
                "Face",
                (x, y-10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 255, 0),
                2
            )
        
        return output
    
    def extract_face_roi(self, image: np.ndarray, face: Tuple[int, int, int, int], 
                        padding: float = 0.2) -> Optional[np.ndarray]:
        """
        Extract face region of interest with padding
        
        Args:
            image: Input image
            face: Face bounding box (x, y, w, h)
            padding: Padding ratio around face
        
        Returns:
            Face ROI image or None
        """
        try:
            x, y, w, h = face
            
            # Add padding
            pad_w = int(w * padding)
            pad_h = int(h * padding)
            
            x1 = max(0, x - pad_w)
            y1 = max(0, y - pad_h)
            x2 = min(image.shape[1], x + w + pad_w)
            y2 = min(image.shape[0], y + h + pad_h)
            
            face_roi = image[y1:y2, x1:x2]
            
            return face_roi if face_roi.size > 0 else None
            
        except Exception as e:
            logger.error(f"Face ROI extraction error: {e}")
            return None