55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
from yoloface import face_analysis
|
|
|
|
face_detector = face_analysis()
|
|
|
|
|
|
def equalize_image(data: np.ndarray):
|
|
data_hsv = cv2.cvtColor(data, cv2.COLOR_RGB2HSV)
|
|
data_hsv[:, :, 2] = cv2.equalizeHist(data_hsv[:, :, 2])
|
|
return cv2.cvtColor(data_hsv, cv2.COLOR_HSV2RGB)
|
|
|
|
|
|
def find_face_bbox_yolo(data: np.ndarray):
|
|
_, box, conf = face_detector.face_detection(frame_arr=data, frame_status=True, model='full')
|
|
if len(box) < 1:
|
|
return None, None
|
|
return box, conf
|
|
|
|
|
|
def find_face_bbox(data: np.ndarray):
|
|
classifier_files = [
|
|
'haarcascades/haarcascade_frontalface_default.xml',
|
|
'haarcascades/haarcascade_frontalface_alt.xml',
|
|
'haarcascades/haarcascade_frontalface_alt2.xml',
|
|
'haarcascades/haarcascade_profileface.xml',
|
|
'haarcascades/haarcascade_glasses.xml',
|
|
'lbpcascade_animeface.xml',
|
|
]
|
|
data_equalized = equalize_image(data)
|
|
data_gray = cv2.cvtColor(data_equalized, cv2.COLOR_RGB2GRAY)
|
|
face_coords, conf = find_face_bbox_yolo(cv2.cvtColor(data_equalized, cv2.COLOR_RGB2BGR))
|
|
if face_coords is not None:
|
|
return face_coords[0]
|
|
|
|
for classifier in classifier_files:
|
|
face_cascade = cv2.CascadeClassifier(classifier)
|
|
face_coords = face_cascade.detectMultiScale(data_gray, 1.1, 3)
|
|
if face_coords is not None:
|
|
break
|
|
return max(face_coords, key=lambda v: v[2]*v[3])
|
|
|
|
|
|
def crop_face(data: np.ndarray, bounding_box) -> np.ndarray:
|
|
x, y, w, h = bounding_box
|
|
# Extending the boxes
|
|
factor = 0.4
|
|
x, y = round(x - factor * w), round(y - factor * h)
|
|
w, h = round(w + factor * w * 2), round(h + factor * h * 2)
|
|
y = max(y, 0)
|
|
x = max(x, 0)
|
|
|
|
face = data[y:y + h, x:x + w]
|
|
return face
|