|
import cv2 |
|
import mediapipe as mp |
|
|
|
mp_face_mesh = mp.solutions.face_mesh |
|
mp_drawing = mp.solutions.drawing_utils |
|
mp_drawing_styles = mp.solutions.drawing_styles |
|
mp_face_detection = mp.solutions.face_detection |
|
|
|
|
|
def face_rect(images): |
|
with mp_face_detection.FaceDetection( |
|
model_selection=1, min_detection_confidence=0.5 |
|
) as face_detection: |
|
for image_cv2 in images: |
|
|
|
results = face_detection.process(cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)) |
|
|
|
|
|
if not results.detections: |
|
yield None |
|
for detection in results.detections: |
|
yield _get_bounding_rect(image_cv2, detection) |
|
|
|
|
|
def _get_bounding_rect( |
|
image: mp_drawing.np.ndarray, |
|
detection: mp_drawing.detection_pb2.Detection, |
|
): |
|
""" |
|
Stolen from mediapipe.solutions.drawing_utils.draw_detection() |
|
""" |
|
if not detection.location_data: |
|
return |
|
if image.shape[2] != mp_drawing._BGR_CHANNELS: |
|
raise ValueError("Input image must contain three channel bgr data.") |
|
image_rows, image_cols, _ = image.shape |
|
|
|
location = detection.location_data |
|
|
|
|
|
if not location.HasField("relative_bounding_box"): |
|
return |
|
relative_bounding_box = location.relative_bounding_box |
|
rect_start_point = mp_drawing._normalized_to_pixel_coordinates( |
|
relative_bounding_box.xmin, relative_bounding_box.ymin, image_cols, image_rows |
|
) |
|
rect_end_point = mp_drawing._normalized_to_pixel_coordinates( |
|
relative_bounding_box.xmin + relative_bounding_box.width, |
|
relative_bounding_box.ymin + relative_bounding_box.height, |
|
image_cols, |
|
image_rows, |
|
) |
|
|
|
return *rect_start_point, *rect_end_point |
|
|
|
|