Merge pull request #845 from tantonj/feature/compreface_subjects

added new functionality to use compreface subjects for facial recognition

docs/src/pages/components-explorer/components/compreface/config.json

@@ -162,6 +162,13 @@
             "description": "If true includes system information like execution_time and plugin_version fields.",
             "optional": true,
             "default": false
+          },
+          {
+            "type": "boolean",
+            "name": "use_subjects",
+            "description": "If true ignores the face_recognition folder structure and uses subjects inside compreface. User can then call the api/v1/compreface/update_subjects endpoint to update entities if new subjects are added into compreface.",
+            "optional": true,
+            "default": false
           }
         ],
         "name": "face_recognition",

viseron/components/compreface/__init__.py

@@ -26,13 +26,15 @@
     CONFIG_SIMILARITTY_THRESHOLD,
     CONFIG_STATUS,
     CONFIG_TRAIN,
+    CONFIG_USE_SUBJECTS,
     DEFAULT_DET_PROB_THRESHOLD,
     DEFAULT_FACE_PLUGINS,
     DEFAULT_LIMIT,
     DEFAULT_PREDICTION_COUNT,
     DEFAULT_SIMILARITTY_THRESHOLD,
     DEFAULT_STATUS,
     DEFAULT_TRAIN,
+    DEFAULT_USE_SUBJECTS,
     DESC_API_KEY,
     DESC_COMPONENT,
     DESC_DET_PROB_THRESHOLD,
@@ -45,6 +47,7 @@
     DESC_SIMILARITY_THRESHOLD,
     DESC_STATUS,
     DESC_TRAIN,
+    DESC_USE_SUBJECTS,
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -82,6 +85,11 @@
         vol.Optional(
             CONFIG_STATUS, default=DEFAULT_STATUS, description=DESC_STATUS
         ): bool,
+        vol.Optional(
+            CONFIG_USE_SUBJECTS,
+            default=DEFAULT_USE_SUBJECTS,
+            description=DESC_USE_SUBJECTS,
+        ): bool,
     }
 )
 
@@ -120,6 +128,6 @@ def setup(vis: Viseron, config) -> bool:
             )
 
         if config[CONFIG_FACE_RECOGNITION][CONFIG_TRAIN]:
-            CompreFaceTrain(config)
+            CompreFaceTrain(vis, config)
 
     return True

viseron/components/compreface/const.py

@@ -3,7 +3,7 @@
 from typing import Final
 
 COMPONENT = "compreface"
-
+SUBJECTS = "subjects"
 
 # CONFIG_SCHEMA constants
 CONFIG_FACE_RECOGNITION = "face_recognition"
@@ -25,6 +25,7 @@
 CONFIG_PREDICTION_COUNT = "prediction_count"
 CONFIG_FACE_PLUGINS = "face_plugins"
 CONFIG_STATUS = "status"
+CONFIG_USE_SUBJECTS = "use_subjects"
 
 DEFAULT_TRAIN = False
 DEFAULT_DET_PROB_THRESHOLD = 0.8
@@ -33,6 +34,7 @@
 DEFAULT_PREDICTION_COUNT = 1
 DEFAULT_FACE_PLUGINS: Final = None
 DEFAULT_STATUS = False
+DEFAULT_USE_SUBJECTS = False
 
 DESC_TRAIN = (
     "Train CompreFace to recognize faces on Viseron start. "
@@ -65,3 +67,8 @@
 DESC_STATUS = (
     "If true includes system information like execution_time and plugin_version fields."
 )
+DESC_USE_SUBJECTS = (
+    "If true ignores the face_recognition folder structure and uses subjects "
+    "inside compreface. User can then call the api/v1/compreface/update_subjects "
+    "endpoint to update entities if new subjects are added into compreface."
+)

viseron/components/compreface/face_recognition.py

@@ -7,12 +7,12 @@
 
 import cv2
 from compreface import CompreFace
-from compreface.collections import FaceCollection
-from compreface.service import RecognitionService
+from compreface.collections import FaceCollection, Subjects
 from face_recognition.face_recognition_cli import image_files_in_folder
 
 from viseron.domains.camera.shared_frames import SharedFrame
 from viseron.domains.face_recognition import AbstractFaceRecognition
+from viseron.domains.face_recognition.binary_sensor import FaceDetectionBinarySensor
 from viseron.domains.face_recognition.const import CONFIG_FACE_RECOGNITION_PATH
 from viseron.helpers import calculate_absolute_coords
 
@@ -28,6 +28,8 @@
     CONFIG_PREDICTION_COUNT,
     CONFIG_SIMILARITTY_THRESHOLD,
     CONFIG_STATUS,
+    CONFIG_USE_SUBJECTS,
+    SUBJECTS,
 )
 
 if TYPE_CHECKING:
@@ -49,7 +51,11 @@ class FaceRecognition(AbstractFaceRecognition):
 
     def __init__(self, vis: Viseron, config, camera_identifier) -> None:
         super().__init__(
-            vis, COMPONENT, config[CONFIG_FACE_RECOGNITION], camera_identifier
+            vis,
+            COMPONENT,
+            config[CONFIG_FACE_RECOGNITION],
+            camera_identifier,
+            not config[CONFIG_FACE_RECOGNITION][CONFIG_USE_SUBJECTS],
         )
 
         options = {
@@ -72,9 +78,28 @@ def __init__(self, vis: Viseron, config, camera_identifier) -> None:
             port=str(config[CONFIG_FACE_RECOGNITION][CONFIG_PORT]),
             options=options,
         )
-        self._recognition: RecognitionService = self._compre_face.init_face_recognition(
+        if COMPONENT not in self._vis.data:
+            self._vis.data[COMPONENT] = {}
+        self._vis.data[COMPONENT][
+            CONFIG_FACE_RECOGNITION
+        ] = self._compre_face.init_face_recognition(
             config[CONFIG_FACE_RECOGNITION][CONFIG_API_KEY]
         )
+        if config[CONFIG_FACE_RECOGNITION][CONFIG_USE_SUBJECTS]:
+            self.update_subject_entities()
+
+    def update_subject_entities(self) -> None:
+        """Update entities with binary face recognition subjects from compreface."""
+        subjects: Subjects = self._vis.data[COMPONENT][
+            CONFIG_FACE_RECOGNITION
+        ].get_subjects()
+        for subject in subjects.list()[SUBJECTS]:
+            binary_sensor = FaceDetectionBinarySensor(self._vis, self._camera, subject)
+            if not self._vis.states.entity_exists(binary_sensor):
+                self._vis.add_entity(
+                    COMPONENT,
+                    FaceDetectionBinarySensor(self._vis, self._camera, subject),
+                )
 
     def face_recognition(
         self, shared_frame: SharedFrame, detected_object: DetectedObject
@@ -93,7 +118,7 @@ def face_recognition(
         cropped_frame = frame[y1:y2, x1:x2].copy()
 
         try:
-            detections = self._recognition.recognize(
+            detections = self._vis.data[COMPONENT][CONFIG_FACE_RECOGNITION].recognize(
                 cv2.imencode(".jpg", cropped_frame)[1].tobytes(),
             )
         except Exception as error:  # pylint: disable=broad-except
@@ -105,7 +130,7 @@ def face_recognition(
             return
 
         for result in detections["result"]:
-            subject = result["subjects"][0]
+            subject = result[SUBJECTS][0]
             if subject["similarity"] >= self._config[CONFIG_SIMILARITTY_THRESHOLD]:
                 self._logger.debug(f"Face found: {subject}")
                 self.known_face_found(
@@ -137,8 +162,9 @@ def face_recognition(
 class CompreFaceTrain:
     """Train CompreFace to recognize faces."""
 
-    def __init__(self, config) -> None:
+    def __init__(self, vis: Viseron, config) -> None:
         self._config = config
+        self._vis = vis
 
         options = {
             CONFIG_LIMIT: config[CONFIG_FACE_RECOGNITION][CONFIG_LIMIT],
@@ -160,10 +186,14 @@ def __init__(self, config) -> None:
             port=str(config[CONFIG_FACE_RECOGNITION][CONFIG_PORT]),
             options=options,
         )
-        self._recognition: RecognitionService = self._compre_face.init_face_recognition(
+        self._vis.data[COMPONENT][
+            CONFIG_FACE_RECOGNITION
+        ] = self._compre_face.init_face_recognition(
             config[CONFIG_FACE_RECOGNITION][CONFIG_API_KEY]
         )
-        self._face_collection: FaceCollection = self._recognition.get_face_collection()
+        self._face_collection: FaceCollection = self._vis.data[COMPONENT][
+            CONFIG_FACE_RECOGNITION
+        ].get_face_collection()
 
         self.train()
 

viseron/components/webserver/api/v1/__init__.py

@@ -3,6 +3,7 @@
 from viseron.components.webserver.api.v1.auth import AuthAPIHandler
 from viseron.components.webserver.api.v1.camera import CameraAPIHandler
 from viseron.components.webserver.api.v1.cameras import CamerasAPIHandler
+from viseron.components.webserver.api.v1.compreface import ComprefaceAPIHandler
 from viseron.components.webserver.api.v1.config import ConfigAPIHandler
 from viseron.components.webserver.api.v1.events import EventsAPIHandler
 from viseron.components.webserver.api.v1.hls import HlsAPIHandler
@@ -13,6 +14,7 @@
     "AuthAPIHandler",
     "CameraAPIHandler",
     "CamerasAPIHandler",
+    "ComprefaceAPIHandler",
     "ConfigAPIHandler",
     "EventsAPIHandler",
     "HlsAPIHandler",

viseron/components/webserver/api/v1/compreface.py

@@ -0,0 +1,58 @@
+"""Config API Handler."""
+import logging
+from http import HTTPStatus
+
+from compreface.collections import Subjects
+
+from viseron.components.compreface.const import (
+    COMPONENT,
+    CONFIG_FACE_RECOGNITION,
+    SUBJECTS,
+)
+from viseron.components.webserver.api.handlers import BaseAPIHandler
+from viseron.const import REGISTERED_DOMAINS
+from viseron.domains.camera.const import DOMAIN as CAMERA_DOMAIN
+from viseron.domains.face_recognition.binary_sensor import FaceDetectionBinarySensor
+
+LOGGER = logging.getLogger(__name__)
+
+
+class ComprefaceAPIHandler(BaseAPIHandler):
+    """Handler for API calls related to compreface."""
+
+    routes = [
+        {
+            "path_pattern": r"/compreface/update_subjects",
+            "supported_methods": ["GET"],
+            "method": "update_subjects",
+        },
+    ]
+
+    async def update_subjects(self) -> None:
+        """Update Viseron subjects."""
+        if COMPONENT not in self._vis.data:
+            self.response_error(
+                status_code=HTTPStatus.BAD_REQUEST,
+                reason="Compreface Recognition not initialized.",
+            )
+        else:
+            subjects: Subjects = self._vis.data[COMPONENT][
+                CONFIG_FACE_RECOGNITION
+            ].get_subjects()
+            added_subjects = []
+            for camera in (
+                self._vis.data[REGISTERED_DOMAINS].get(CAMERA_DOMAIN, {}).values()
+            ):
+                for subject in subjects.list()[SUBJECTS]:
+                    binary_sensor = FaceDetectionBinarySensor(
+                        self._vis, camera, subject
+                    )
+                    if not self._vis.states.entity_exists(binary_sensor):
+                        added_subjects.append(f"{camera.identifier}_{subject}")
+                        self._vis.add_entity(
+                            COMPONENT,
+                            FaceDetectionBinarySensor(self._vis, camera, subject),
+                        )
+            response = {}
+            response["added_subjects"] = added_subjects
+            self.response_success(response=response)

viseron/domains/face_recognition/__init__.py

@@ -112,16 +112,18 @@ def as_dict(self) -> dict[str, Any]:
 class AbstractFaceRecognition(AbstractPostProcessor):
     """Abstract face recognition."""
 
-    def __init__(self, vis, component, config, camera_identifier) -> None:
+    def __init__(
+        self, vis, component, config, camera_identifier, generate_entities=True
+    ) -> None:
         super().__init__(vis, config, camera_identifier)
         self._faces: dict[str, FaceDict] = {}
-
-        for face_dir in os.listdir(config[CONFIG_FACE_RECOGNITION_PATH]):
-            if face_dir == "unknown":
-                continue
-            vis.add_entity(
-                component, FaceDetectionBinarySensor(vis, self._camera, face_dir)
-            )
+        if generate_entities:
+            for face_dir in os.listdir(config[CONFIG_FACE_RECOGNITION_PATH]):
+                if face_dir == "unknown":
+                    continue
+                vis.add_entity(
+                    component, FaceDetectionBinarySensor(vis, self._camera, face_dir)
+                )
 
     @abstractmethod
     def face_recognition(

viseron/states.py

@@ -168,6 +168,10 @@ def _assign_object_id(entity: Entity) -> None:
         else:
             entity.object_id = slugify(entity.name)
 
+    def entity_exists(self, entity: Entity) -> bool:
+        """Return if entity has already been added."""
+        return self._generate_entity_id(entity) in self._registry
+
     def _generate_entity_id(self, entity: Entity) -> str:
         """Generate entity id for an entity."""
         self._assign_object_id(entity)