src/face_detector.py

from email.mime import image

import cv2
import numpy as np
from facenet_pytorch import MTCNN
from src.config_reader import Config

class FaceDetector:
    def __init__(self):
        self.cfg = Config()
        self.detector = MTCNN(
        keep_all=False,
        device='cpu'
    )
    
    def detect(self, image):
    # facenet_pytorch MTCNN direct
        boxes, probs, landmarks = self.detector.detect(image, landmarks=True)
    
        if boxes is None:
            return None
    
    # First face return
        return {
            'box': [int(boxes[0][0]), int(boxes[0][1]), 
                    int(boxes[0][2]-boxes[0][0]), 
                    int(boxes[0][3]-boxes[0][1])],
            'keypoints': {
                'left_eye': (int(landmarks[0][0][0]), int(landmarks[0][0][1])),
                'right_eye': (int(landmarks[0][1][0]), int(landmarks[0][1][1])),
        }
    }
    
    def align(self, image, keypoints):
        left_eye = keypoints['left_eye']
        right_eye = keypoints['right_eye']  # Implementation for face alignment goes here
        
        dy = right_eye[1] - left_eye[1]
        dx = right_eye[0] - left_eye[0]
        angle = np.degrees(np.arctan2(dy, dx))
        
        center = (image.shape[1] // 2, image.shape[0] // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        aligned = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
        return aligned
        

    def extract(self, image):
        face_data = self.detect(image)
    
        if face_data is None:
            return None
    
        x, y, w, h = face_data['box']
        keypoints = face_data['keypoints']
    
        aligned = self.align(image, keypoints)
    
        pad = 20
        x1 = max(0, x - pad)
        y1 = max(0, y - pad)
        x2 = min(image.shape[1], x + w + pad)
        y2 = min(image.shape[0], y + h + pad)
    
        cropped = aligned[y1:y2, x1:x2]
    
        resized = cv2.resize(cropped, (224, 224))
    
        return resized