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
from typing import List
from typing import List, Union
# 3rd party dependencies
import numpy as np
@ -13,64 +13,67 @@ class RetinaFaceClient(Detector):
def __init__(self):
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:
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:
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):
return resp
for img in imgs:
resp = []
obj = rf.detect_faces(img, model=self.model, threshold=0.9)
for face_idx in obj.keys():
identity = obj[face_idx]
detection = identity["facial_area"]
if isinstance(obj, dict):
for face_idx in obj.keys():
identity = obj[face_idx]
detection = identity["facial_area"]
y = detection[1]
h = detection[3] - y
x = detection[0]
w = detection[2] - x
y = detection[1]
h = detection[3] - y
x = detection[0]
w = detection[2] - x
# retinaface sets left and right eyes with respect to the person
left_eye = identity["landmarks"]["left_eye"]
right_eye = identity["landmarks"]["right_eye"]
nose = identity["landmarks"].get("nose")
mouth_right = identity["landmarks"].get("mouth_right")
mouth_left = identity["landmarks"].get("mouth_left")
left_eye = tuple(int(i) for i in identity["landmarks"]["left_eye"])
right_eye = tuple(int(i) for i in identity["landmarks"]["right_eye"])
nose = identity["landmarks"].get("nose")
mouth_right = identity["landmarks"].get("mouth_right")
mouth_left = identity["landmarks"].get("mouth_left")
# eyes are list of float, need to cast them tuple of int
left_eye = tuple(int(i) for i in left_eye)
right_eye = tuple(int(i) for i in right_eye)
if nose is not None:
nose = tuple(int(i) for i in nose)
if mouth_right is not None:
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)
if nose is not None:
nose = tuple(int(i) for i in nose)
if mouth_right is not None:
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(
x=x,
y=y,
w=w,
h=h,
left_eye=left_eye,
right_eye=right_eye,
confidence=confidence,
nose=nose,
mouth_left=mouth_left,
mouth_right=mouth_right,
)
facial_area = FacialAreaRegion(
x=x,
y=y,
w=w,
h=h,
left_eye=left_eye,
right_eye=right_eye,
confidence=confidence,
nose=nose,
mouth_left=mouth_left,
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]