From 544df18982af60b301bf5f7d50a9087518bf64f1 Mon Sep 17 00:00:00 2001 From: Hugo Burton Date: Fri, 23 Aug 2024 10:20:50 +1000 Subject: [PATCH] Simplify detect faces code --- .../face_landmark_estimator.py | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/eye_tracking/gaze/head_pose_estimation/face_landmark_estimator.py b/eye_tracking/gaze/head_pose_estimation/face_landmark_estimator.py index dbdc6040..a6a56608 100644 --- a/eye_tracking/gaze/head_pose_estimation/face_landmark_estimator.py +++ b/eye_tracking/gaze/head_pose_estimation/face_landmark_estimator.py @@ -1,7 +1,5 @@ from typing import List -import face_alignment -import face_alignment.detection.sfd import mediapipe import numpy as np from omegaconf import DictConfig @@ -19,21 +17,14 @@ def __init__(self, config: DictConfig): ) def detect_faces(self, image: np.ndarray) -> List[Face]: - return self._detect_faces_mediapipe(image) - - def detect_faces_raw(self, image: np.ndarray) -> List[np.ndarray]: - if self.mode == "mediapipe": - return self._detect_faces_mediapipe_raw(image) - else: - raise ValueError - - def _detect_faces_mediapipe(self, image: np.ndarray) -> List[Face]: """ Calculated landmarks scaled to the image size with a bounding box + :param image: RGB image + :return: List of faces """ h, w = image.shape[:2] - faces_landmarks = self._detect_faces_mediapipe_raw(image) + faces_landmarks = self._detect_faces_raw(image) detected = [] if faces_landmarks: for face in faces_landmarks: @@ -43,9 +34,17 @@ def _detect_faces_mediapipe(self, image: np.ndarray) -> List[Face]: detected.append(Face(bbox, pts)) return detected - def _detect_faces_mediapipe_raw(self, image: np.ndarray) -> List[np.ndarray]: + def detect_faces_raw(self, image: np.ndarray) -> List[np.ndarray]: + if self.mode == "mediapipe": + return self._detect_faces_raw(image) + else: + raise ValueError + + def _detect_faces_raw(self, image: np.ndarray) -> List[np.ndarray]: """ - Returns landmarks as they come from the mediapipe model + Returns landmarks as they come from the mediapipe model (not scaled to the image size) + :param image: RGB image + :return: List of faces landmarks """ predictions = self.detector.process(self._get_bgr_frame(image)) faces_landmarks = []