Upgrade to new arcface model

This commit is contained in:
Nicolas Mowen 2025-03-25 14:04:41 -06:00
parent 2ab934a9d2
commit cc03075cbe
3 changed files with 16 additions and 13 deletions

View File

@ -44,6 +44,7 @@ class FaceRecognizer(ABC):
output_height: int,
) -> np.ndarray:
# landmark is run on grayscale images
if image.ndim == 3:
land_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
@ -230,7 +231,6 @@ class FaceNetRecognizer(FaceRecognizer):
if img is None:
continue
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = self.align_face(img, img.shape[1], img.shape[0])
emb = self.face_embedder([img])[0].squeeze()
face_embeddings_map[name].append(emb)
@ -241,8 +241,7 @@ class FaceNetRecognizer(FaceRecognizer):
return
for name, embs in face_embeddings_map.items():
norms = np.linalg.norm(embs, axis=-1, keepdims=True)
self.mean_embs[name] = stats.trim_mean(embs / norms, 0.15)
self.mean_embs[name] = stats.trim_mean(embs, 0.15)
def classify(self, face_image):
if not self.landmark_detector:
@ -255,14 +254,13 @@ class FaceNetRecognizer(FaceRecognizer):
return None
# face recognition is best run on grayscale images
img = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)
# get blur factor before aligning face
blur_factor = self.get_blur_factor(img)
blur_factor = self.get_blur_factor(face_image)
logger.debug(f"face detected with bluriness {blur_factor}")
# align face and run recognition
img = self.align_face(img, img.shape[1], img.shape[0])
img = self.align_face(face_image, face_image.shape[1], face_image.shape[0])
embedding = self.face_embedder([img])[0].squeeze()
score = 0

View File

@ -85,6 +85,7 @@ class BaseEmbedding(ABC):
input_names = self.runner.get_input_names()
onnx_inputs = {name: [] for name in input_names}
input: dict[str, any]
print(f"onnx inputs are {input_names}")
for input in processed:
for key, value in input.items():
if key in input_names:

View File

@ -13,7 +13,7 @@ from .runner import ONNXModelRunner
logger = logging.getLogger(__name__)
FACE_EMBEDDING_SIZE = 160
FACE_EMBEDDING_SIZE = 112
class FaceNetEmbedding(BaseEmbedding):
@ -77,9 +77,9 @@ class FaceNetEmbedding(BaseEmbedding):
og = np.array(pil).astype(np.float32)
# Image must be FACE_EMBEDDING_SIZExFACE_EMBEDDING_SIZE
og_h, og_w = og.shape
og_h, og_w, channels = og.shape
frame = np.zeros(
(FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, 3), dtype=np.float32
(FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, channels), dtype=np.float32
)
# compute center offset
@ -87,8 +87,12 @@ class FaceNetEmbedding(BaseEmbedding):
y_center = (FACE_EMBEDDING_SIZE - og_h) // 2
# copy img image into center of result image
frame[y_center : y_center + og_h, x_center : x_center + og_w, 0] = og
frame[y_center : y_center + og_h, x_center : x_center + og_w, 1] = og
frame[y_center : y_center + og_h, x_center : x_center + og_w, 2] = og
frame[y_center : y_center + og_h, x_center : x_center + og_w] = og
# run arcface normalization
normalized_image = frame.astype(np.float32) / 255.0
frame = (normalized_image - 0.5) / 0.5
frame = np.transpose(frame, (2, 0, 1))
frame = np.expand_dims(frame, axis=0)
return [{"input_2": frame}]
return [{"data": frame}]