from facenet_pytorch import MTCNN
import torch
from PIL import Image
import numpy as np


class FaceDetector:
    def __init__(self):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = MTCNN(keep_all=True, post_process=False, device=device)

    def detect_bboxes(self, img: Image.Image) -> np.ndarray:
        boxes, probs = self.model.detect(img)
        return boxes
    
    def extract_faces(self, img: Image.Image, bboxes) -> torch.Tensor:
        face_images_tensors = self.model.extract(img, bboxes, save_path='dev/extracted/extracted.jpg')
        return face_images_tensors