55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
from mtcnn import MTCNN
|
|
from facenet_pytorch import MTCNN
|
|
import torch
|
|
import math
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from Face_facial.face_detector import YoloDetector
|
|
|
|
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.model = YoloDetector(target_size=720, device="cuda:0", min_face=10)
|
|
|
|
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):
|
|
boxes, landmarks = self.model.predict(frame)
|
|
|
|
if len(boxes) > 0:
|
|
|
|
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) >= 3:
|
|
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) > 65):
|
|
return True, True, [x1 - 20, y1 - 20, x2 + 20, y2 + 20]
|
|
|
|
return True, False, []
|
|
|
|
return False, False, []
|