Lazy-load face detector and embedder

frame_extraction/src/frame_extraction/face.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Iterable, List, Tuple
+from dataclasses import dataclass, field
+from typing import Iterable, List
 
 import numpy as np
 import torch
@@ -10,18 +9,22 @@ from facenet_pytorch import InceptionResnetV1, MTCNN
 from PIL import Image
 
 
-@dataclass(slots=True)
+@dataclass
 class FaceDetector:
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    device: str = field(default_factory=lambda: "cuda" if torch.cuda.is_available() else "cpu")
     min_face_size: int = 60
-
-    def __post_init__(self) -> None:
-        self.model = MTCNN(
-            keep_all=True,
-            device=self.device,
-            min_face_size=self.min_face_size,
-            post_process=False,
-        )
+    _model: MTCNN | None = field(init=False, default=None, repr=False)
+
+    @property
+    def model(self) -> MTCNN:
+        if self._model is None:
+            self._model = MTCNN(
+                keep_all=True,
+                device=self.device,
+                min_face_size=self.min_face_size,
+                post_process=False,
+            )
+        return self._model
 
     def detect(self, image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
         pil = Image.fromarray(cv2_to_rgb(image))
@@ -31,13 +34,17 @@ class FaceDetector:
         return boxes.astype(np.float32), probs.astype(np.float32)
 
 
-@dataclass(slots=True)
+@dataclass
 class FaceEmbedder:
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    device: str = field(default_factory=lambda: "cuda" if torch.cuda.is_available() else "cpu")
     batch_size: int = 16
+    _model: InceptionResnetV1 | None = field(init=False, default=None, repr=False)
 
-    def __post_init__(self) -> None:
-        self.model = InceptionResnetV1(pretrained="vggface2").eval().to(self.device)
+    @property
+    def model(self) -> InceptionResnetV1:
+        if self._model is None:
+            self._model = InceptionResnetV1(pretrained="vggface2").eval().to(self.device)
+        return self._model
 
     @torch.no_grad()
     def embed(self, crops: Iterable[Image.Image]) -> np.ndarray: