|
import torch
|
from typing import Optional, List
|
import numpy
|
from insightface.app import FaceAnalysis
|
from insightface.app.common import Face
|
import numpy as np
|
|
SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个
|
|
|
class FaceAnalyser:
|
def __init__(self):
|
self.device = self._select_device()
|
self._analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider' if self.device == 'cuda' else 'CPUExecutionProvider'])
|
self._analyser.prepare(ctx_id=0, det_size=(640, 640))
|
|
def _select_device(self):
|
if torch.cuda.is_available():
|
return 'cuda'
|
elif torch.backends.mps.is_available():
|
return 'mps'
|
else:
|
return 'cpu'
|
|
def get_many_faces(self,frame: np.ndarray) -> Optional[List[Face]]:
|
try:
|
return self._analyser.get(frame)
|
except Exception as e:
|
print(f"{e}")
|
return None
|
|
def get_one_face(self,frame: np.ndarray, position: int = 0) -> Optional[Face]:
|
many_faces = self.get_many_faces(frame)
|
if many_faces:
|
try:
|
return many_faces[position]
|
except IndexError:
|
return many_faces[-1]
|
return None
|
|
|
def find_similar_face(self,frame: np.ndarray, reference_face: Face) -> Optional[Face]:
|
many_faces = self.get_many_faces(frame)
|
if many_faces:
|
for face in many_faces:
|
if hasattr(face, 'normed_embedding') and hasattr(reference_face, 'normed_embedding'):
|
distance = numpy.sum(numpy.square(face.normed_embedding - reference_face.normed_embedding))
|
if distance < SIMILAR_FACE_DISTANCE:
|
return face
|
return None
|
