from facenet_pytorch import MTCNN
from PIL import Image
import torch
import gradio as gr

mtcnn = MTCNN(margin=0, thresholds=[0.85, 0.95, 0.95], device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))

def extract_faces(frames, progress=gr.Progress(track_tqdm=True)):
    face_crops = []
    for frame in progress.tqdm(frames, desc='Detecting faces'):
        img = Image.fromarray(frame)

        boxes, _ = mtcnn.detect(img)

        if boxes is not None:
            for i, box in enumerate(boxes):
                x1, y1, x2, y2 = [int(b) for b in box]
                face = img.crop((x1, y1, x2, y2))
                face_crops.append(face)
    return face_crops