Spaces:
Sleeping
Sleeping
File size: 3,197 Bytes
7a77962 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import cv2
import torch
import numpy as np
from PIL import Image
try:
from facenet_pytorch import MTCNN
except ImportError:
print("facenet_pytorch not installed. Please install it using: pip install facenet-pytorch")
MTCNN = None
class FaceDetector:
def __init__(self, device='cuda', image_size=224, margin=0):
self.device = device
if MTCNN is not None:
self.mtcnn = MTCNN(
image_size=image_size,
margin=margin,
keep_all=True, # We'll select the largest ourselves
device=device,
post_process=False # We want the raw image, not normalized
)
else:
self.mtcnn = None
def process_video(self, video_path, fps=5):
"""
Extract faces from video.
Returns a list of (face_rgb, frame_idx) tuples.
"""
if self.mtcnn is None:
raise ImportError("MTCNN not initialized.")
cap = cv2.VideoCapture(video_path)
video_fps = cap.get(cv2.CAP_PROP_FPS)
frame_interval = int(max(1, round(video_fps / fps)))
faces = []
frame_idx = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
if frame_idx % frame_interval == 0:
# Convert BGR to RGB
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(frame_rgb)
# Detect faces
boxes, _ = self.mtcnn.detect(pil_img)
if boxes is not None and len(boxes) > 0:
areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in boxes]
largest_idx = np.argmax(areas)
box = boxes[largest_idx]
b = [int(max(0, c)) for c in box]
face_img = frame_rgb[b[1]:b[3], b[0]:b[2]]
if face_img.size > 0:
face_img = cv2.resize(face_img, (224, 224))
faces.append(face_img)
frame_idx += 1
cap.release()
return faces
def process_image(self, image_path):
"""
Process a single image path.
"""
if self.mtcnn is None:
raise ImportError("MTCNN not initialized.")
frame = cv2.imread(image_path)
if frame is None:
return None
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(frame_rgb)
boxes, _ = self.mtcnn.detect(pil_img)
if boxes is not None and len(boxes) > 0:
areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in boxes]
largest_idx = np.argmax(areas)
box = boxes[largest_idx]
b = [int(max(0, c)) for c in box]
face_img = frame_rgb[b[1]:b[3], b[0]:b[2]]
if face_img.size > 0:
face_img = cv2.resize(face_img, (224, 224))
return face_img
return None
|