Update writing images

This commit is contained in:
Nicolas Mowen 2024-11-21 17:15:02 -07:00
parent 646d878ccf
commit c4e944dc93
2 changed files with 26 additions and 14 deletions

View File

@ -3,8 +3,6 @@
import base64 import base64
import logging import logging
import os import os
import random
import string
import time import time
from numpy import ndarray from numpy import ndarray
@ -14,7 +12,6 @@ from frigate.comms.inter_process import InterProcessRequestor
from frigate.config import FrigateConfig from frigate.config import FrigateConfig
from frigate.const import ( from frigate.const import (
CONFIG_DIR, CONFIG_DIR,
FACE_DIR,
UPDATE_EMBEDDINGS_REINDEX_PROGRESS, UPDATE_EMBEDDINGS_REINDEX_PROGRESS,
UPDATE_MODEL_STATE, UPDATE_MODEL_STATE,
) )

View File

@ -3,7 +3,9 @@
import base64 import base64
import logging import logging
import os import os
import random
import re import re
import string
import threading import threading
from multiprocessing.synchronize import Event as MpEvent from multiprocessing.synchronize import Event as MpEvent
from typing import Optional from typing import Optional
@ -22,7 +24,12 @@ from frigate.comms.event_metadata_updater import (
from frigate.comms.events_updater import EventEndSubscriber, EventUpdateSubscriber from frigate.comms.events_updater import EventEndSubscriber, EventUpdateSubscriber
from frigate.comms.inter_process import InterProcessRequestor from frigate.comms.inter_process import InterProcessRequestor
from frigate.config import FrigateConfig from frigate.config import FrigateConfig
from frigate.const import CLIPS_DIR, FRIGATE_LOCALHOST, UPDATE_EVENT_DESCRIPTION from frigate.const import (
CLIPS_DIR,
FACE_DIR,
FRIGATE_LOCALHOST,
UPDATE_EVENT_DESCRIPTION,
)
from frigate.embeddings.lpr.lpr import LicensePlateRecognition from frigate.embeddings.lpr.lpr import LicensePlateRecognition
from frigate.events.types import EventTypeEnum from frigate.events.types import EventTypeEnum
from frigate.genai import get_genai_client from frigate.genai import get_genai_client
@ -146,12 +153,14 @@ class EmbeddingMaintainer(threading.Thread):
if not self.face_recognition_enabled: if not self.face_recognition_enabled:
return False return False
if data.get("cropped"): rand_id = "".join(
self.embeddings.embed_face( random.choices(string.ascii_lowercase + string.digits, k=6)
data["face_name"],
base64.b64decode(data["image"]),
upsert=True,
) )
label = data["face_name"]
id = f"{label}-{rand_id}"
if data.get("cropped"):
pass
else: else:
img = cv2.imdecode( img = cv2.imdecode(
np.frombuffer( np.frombuffer(
@ -165,12 +174,18 @@ class EmbeddingMaintainer(threading.Thread):
return False return False
face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]] face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]]
ret, webp = cv2.imencode( ret, thumbnail = cv2.imencode(
".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100] ".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100]
) )
self.embeddings.embed_face(
data["face_name"], webp.tobytes(), upsert=True # write face to library
) folder = os.path.join(FACE_DIR, label)
file = os.path.join(folder, f"{id}.webp")
os.makedirs(folder, exist_ok=True)
# save face image
with open(file, "wb") as output:
output.write(thumbnail.tobytes())
self.face_classifier.clear_classifier() self.face_classifier.clear_classifier()
return True return True
@ -473,7 +488,7 @@ class EmbeddingMaintainer(threading.Thread):
json={ json={
"camera": obj_data.get("camera"), "camera": obj_data.get("camera"),
"subLabel": sub_label, "subLabel": sub_label,
"subLabelScore": score "subLabelScore": score,
}, },
) )