File size: 702 Bytes
0077a91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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
|