pseudo batched retinaface

This commit is contained in:
galthran-wq 2025-02-16 13:54:28 +00:00
parent 619930cd1e
commit b544a2d866

View File

@ -1,5 +1,5 @@
# built-in dependencies # built-in dependencies
from typing import List from typing import List, Union
# 3rd party dependencies # 3rd party dependencies
import numpy as np import numpy as np
@ -13,64 +13,67 @@ class RetinaFaceClient(Detector):
def __init__(self): def __init__(self):
self.model = rf.build_model() self.model = rf.build_model()
def detect_faces(self, img: np.ndarray) -> List[FacialAreaRegion]: def detect_faces(self, img: Union[np.ndarray, List[np.ndarray]]) -> Union[List[FacialAreaRegion], List[List[FacialAreaRegion]]]:
""" """
Detect and align face with retinaface Detect and align faces with retinaface in an image or a list of images
Args: Args:
img (np.ndarray): pre-loaded image as numpy array img (Union[np.ndarray, List[np.ndarray]]): pre-loaded image as numpy array or a list of those
Returns: Returns:
results (List[FacialAreaRegion]): A list of FacialAreaRegion objects results (Union[List[FacialAreaRegion], List[List[FacialAreaRegion]]]): A list or a list of lists of FacialAreaRegion objects
""" """
resp = [] if isinstance(img, np.ndarray):
imgs = [img]
else:
imgs = img
obj = rf.detect_faces(img, model=self.model, threshold=0.9) batch_results = []
if not isinstance(obj, dict): for img in imgs:
return resp resp = []
obj = rf.detect_faces(img, model=self.model, threshold=0.9)
for face_idx in obj.keys(): if isinstance(obj, dict):
identity = obj[face_idx] for face_idx in obj.keys():
detection = identity["facial_area"] identity = obj[face_idx]
detection = identity["facial_area"]
y = detection[1] y = detection[1]
h = detection[3] - y h = detection[3] - y
x = detection[0] x = detection[0]
w = detection[2] - x w = detection[2] - x
# retinaface sets left and right eyes with respect to the person left_eye = tuple(int(i) for i in identity["landmarks"]["left_eye"])
left_eye = identity["landmarks"]["left_eye"] right_eye = tuple(int(i) for i in identity["landmarks"]["right_eye"])
right_eye = identity["landmarks"]["right_eye"] nose = identity["landmarks"].get("nose")
nose = identity["landmarks"].get("nose") mouth_right = identity["landmarks"].get("mouth_right")
mouth_right = identity["landmarks"].get("mouth_right") mouth_left = identity["landmarks"].get("mouth_left")
mouth_left = identity["landmarks"].get("mouth_left")
# eyes are list of float, need to cast them tuple of int if nose is not None:
left_eye = tuple(int(i) for i in left_eye) nose = tuple(int(i) for i in nose)
right_eye = tuple(int(i) for i in right_eye) if mouth_right is not None:
if nose is not None: mouth_right = tuple(int(i) for i in mouth_right)
nose = tuple(int(i) for i in nose) if mouth_left is not None:
if mouth_right is not None: mouth_left = tuple(int(i) for i in mouth_left)
mouth_right = tuple(int(i) for i in mouth_right)
if mouth_left is not None:
mouth_left = tuple(int(i) for i in mouth_left)
confidence = identity["score"] confidence = identity["score"]
facial_area = FacialAreaRegion( facial_area = FacialAreaRegion(
x=x, x=x,
y=y, y=y,
w=w, w=w,
h=h, h=h,
left_eye=left_eye, left_eye=left_eye,
right_eye=right_eye, right_eye=right_eye,
confidence=confidence, confidence=confidence,
nose=nose, nose=nose,
mouth_left=mouth_left, mouth_left=mouth_left,
mouth_right=mouth_right, mouth_right=mouth_right,
) )
resp.append(facial_area) resp.append(facial_area)
return resp batch_results.append(resp)
return batch_results if len(batch_results) > 1 else batch_results[0]