| 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 | |