Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

added new functionality to use compreface subjects for facial recognition #845

Merged
merged 4 commits into from
Dec 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
10 changes: 9 additions & 1 deletion viseron/components/compreface/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -45,6 +47,7 @@
DESC_SIMILARITY_THRESHOLD,
DESC_STATUS,
DESC_TRAIN,
DESC_USE_SUBJECTS,
)

LOGGER = logging.getLogger(__name__)
Expand Down Expand Up @@ -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,
}
)

Expand Down Expand Up @@ -120,6 +128,6 @@ def setup(vis: Viseron, config) -> bool:
)

if config[CONFIG_FACE_RECOGNITION][CONFIG_TRAIN]:
CompreFaceTrain(config)
CompreFaceTrain(vis, config)

return True
9 changes: 8 additions & 1 deletion viseron/components/compreface/const.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from typing import Final

COMPONENT = "compreface"

SUBJECTS = "subjects"

# CONFIG_SCHEMA constants
CONFIG_FACE_RECOGNITION = "face_recognition"
Expand All @@ -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
Expand All @@ -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. "
Expand Down Expand Up @@ -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."
)
48 changes: 39 additions & 9 deletions viseron/components/compreface/face_recognition.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -28,6 +28,8 @@
CONFIG_PREDICTION_COUNT,
CONFIG_SIMILARITTY_THRESHOLD,
CONFIG_STATUS,
CONFIG_USE_SUBJECTS,
SUBJECTS,
)

if TYPE_CHECKING:
Expand All @@ -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 = {
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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],
Expand All @@ -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()

Expand Down
2 changes: 2 additions & 0 deletions viseron/components/webserver/api/v1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -13,6 +14,7 @@
"AuthAPIHandler",
"CameraAPIHandler",
"CamerasAPIHandler",
"ComprefaceAPIHandler",
"ConfigAPIHandler",
"EventsAPIHandler",
"HlsAPIHandler",
Expand Down
58 changes: 58 additions & 0 deletions viseron/components/webserver/api/v1/compreface.py
Original file line number Diff line number Diff line change
@@ -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)
18 changes: 10 additions & 8 deletions viseron/domains/face_recognition/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions viseron/states.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading