Increase requirements for face to be set

This commit is contained in:
Nicolas Mowen 2024-10-23 08:08:43 -06:00
parent b7f9eb2ee6
commit b34a51903f

View File

@ -32,6 +32,7 @@ from .embeddings import Embeddings
logger = logging.getLogger(__name__)
REQUIRED_FACES = 2
MAX_THUMBNAILS = 10
@ -65,6 +66,7 @@ class EmbeddingMaintainer(threading.Thread):
self.config.semantic_search.face_recognition.enabled
)
self.requires_face_detection = "face" not in self.config.model.all_attributes
self.detected_faces: dict[str, float] = {}
# create communication for updating event descriptions
self.requestor = InterProcessRequestor()
@ -276,25 +278,28 @@ class EmbeddingMaintainer(threading.Thread):
def _search_face(self, query_embedding: bytes) -> list:
"""Search for the face most closely matching the embedding."""
sql_query = """
sql_query = f"""
SELECT
id,
distance
FROM vec_faces
WHERE face_embedding MATCH ?
AND k = 10 ORDER BY distance
AND k = {REQUIRED_FACES} ORDER BY distance
"""
return self.embeddings.db.execute_sql(sql_query, [query_embedding]).fetchall()
def _process_face(self, obj_data: dict[str, any], frame: np.ndarray) -> None:
"""Look for faces in image."""
id = obj_data["id"]
# don't run for non person objects
if obj_data.get("label") != "person":
logger.debug("Not a processing face for non person object.")
return
# don't overwrite sub label for objects that have one
if obj_data.get("sub_label"):
# don't overwrite sub label for objects that have a sub label
# that is not a face
if obj_data.get("sub_label") and id not in self.detected_faces:
logger.debug(
f"Not processing face due to existing sub label: {obj_data.get('sub_label')}."
)
@ -348,18 +353,25 @@ class EmbeddingMaintainer(threading.Thread):
best_faces = self._search_face(query_embedding)
logger.debug(f"Detected best faces for person as: {best_faces}")
if not best_faces:
if not best_faces or len(best_faces) < REQUIRED_FACES:
return
sub_label = str(best_faces[0][0]).split("-")[0]
score = 1.0 - best_faces[0][1]
avg_score = 0
if score < self.config.semantic_search.face_recognition.threshold:
return None
for face in best_faces:
score = 1.0 - face[1]
if score < self.config.semantic_search.face_recognition.threshold:
return None
avg_score += score
avg_score = avg_score / REQUIRED_FACES
self.detected_faces[id] = avg_score
requests.post(
f"{FRIGATE_LOCALHOST}/api/events/{obj_data['id']}/sub_label",
json={"subLabel": sub_label, "subLabelScore": score},
f"{FRIGATE_LOCALHOST}/api/events/{id}/sub_label",
json={"subLabel": sub_label, "subLabelScore": avg_score},
)
def _create_thumbnail(self, yuv_frame, box, height=500) -> Optional[bytes]: