49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
|
from mtcnn import MTCNN
|
||
|
|
from facenet_pytorch import MTCNN
|
||
|
|
import torch
|
||
|
|
import math
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
K_MULTIPLIER = 1.2
|
||
|
|
|
||
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
|
detector = MTCNN(keep_all=True, device=device)
|
||
|
|
|
||
|
|
|
||
|
|
class FaceDetection:
|
||
|
|
def __init__(self,):
|
||
|
|
self.detector = detector
|
||
|
|
|
||
|
|
def calculate_distance(self, p1, p2) -> float:
|
||
|
|
return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
|
||
|
|
|
||
|
|
def calculate_dis_to_cp(self, cx, cy, face_cx, face_cy) -> float:
|
||
|
|
return math.sqrt((face_cx - cx) ** 2 + (face_cy - cy) ** 2)
|
||
|
|
|
||
|
|
def detect_face(self, frame, cx, cy) -> bool:
|
||
|
|
boxes, probs, landmarks = self.detector.detect(frame, landmarks=True)
|
||
|
|
|
||
|
|
for box, landmark in zip(boxes, landmarks):
|
||
|
|
# Draw bounding box
|
||
|
|
x1, y1, x2, y2 = map(int, box)
|
||
|
|
|
||
|
|
face_cx = int(x1 + (x2 - x1) / 2)
|
||
|
|
face_cy = int(y1 + (y2 - y1) / 2)
|
||
|
|
|
||
|
|
if len(landmark) >= 5:
|
||
|
|
nose = landmark[2]
|
||
|
|
left_eye = landmark[0]
|
||
|
|
right_eye = landmark[1]
|
||
|
|
|
||
|
|
# Calculate distances
|
||
|
|
distance_left = self.calculate_distance(nose, left_eye)
|
||
|
|
distance_right = self.calculate_distance(nose, right_eye)
|
||
|
|
|
||
|
|
# Check if distances exceed threshold
|
||
|
|
if not (distance_left > K_MULTIPLIER * distance_right or distance_right > K_MULTIPLIER * distance_left or
|
||
|
|
self.calculate_dis_to_cp(cx, cy, face_cx, face_cy) > 30):
|
||
|
|
return True
|
||
|
|
|
||
|
|
return False
|