Don't assume landmark file is downloaded

This commit is contained in:
Nicolas Mowen 2025-03-26 07:59:31 -06:00
parent e3d4b84803
commit 2b1fc069eb

View File

@ -18,10 +18,7 @@ class FaceRecognizer(ABC):
def __init__(self, config: FrigateConfig) -> None: def __init__(self, config: FrigateConfig) -> None:
self.config = config self.config = config
self.landmark_detector = cv2.face.createFacemarkLBF() self.init_landmark_detector()
self.landmark_detector.loadModel(
os.path.join(MODEL_CACHE_DIR, "facedet/landmarkdet.yaml")
)
@abstractmethod @abstractmethod
def build(self) -> None: def build(self) -> None:
@ -37,6 +34,13 @@ class FaceRecognizer(ABC):
def classify(self, face_image: np.ndarray) -> tuple[str, float] | None: def classify(self, face_image: np.ndarray) -> tuple[str, float] | None:
pass pass
def init_landmark_detector(self) -> None:
landmark_model = os.path.join(MODEL_CACHE_DIR, "facedet/landmarkdet.yaml")
if os.path.exists(landmark_model):
self.landmark_detector = cv2.face.createFacemarkLBF()
self.landmark_detector.loadModel(landmark_model)
def align_face( def align_face(
self, self,
image: np.ndarray, image: np.ndarray,
@ -130,6 +134,7 @@ class LBPHRecognizer(FaceRecognizer):
def build(self): def build(self):
if not self.landmark_detector: if not self.landmark_detector:
self.init_landmark_detector()
return None return None
labels = [] labels = []
@ -207,6 +212,7 @@ class ArcFaceRecognizer(FaceRecognizer):
def build(self): def build(self):
if not self.landmark_detector: if not self.landmark_detector:
self.init_landmark_detector()
return None return None
face_embeddings_map: dict[str, list[np.ndarray]] = {} face_embeddings_map: dict[str, list[np.ndarray]] = {}