mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-01-31 00:18:55 +01:00
Face recognition improvements (#16034)
This commit is contained in:
parent
9a2de78fc9
commit
0ec536a4e5
@ -24,14 +24,18 @@ def get_faces():
|
|||||||
face_dict: dict[str, list[str]] = {}
|
face_dict: dict[str, list[str]] = {}
|
||||||
|
|
||||||
for name in os.listdir(FACE_DIR):
|
for name in os.listdir(FACE_DIR):
|
||||||
face_dict[name] = []
|
|
||||||
|
|
||||||
face_dir = os.path.join(FACE_DIR, name)
|
face_dir = os.path.join(FACE_DIR, name)
|
||||||
|
|
||||||
if not os.path.isdir(face_dir):
|
if not os.path.isdir(face_dir):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
for file in os.listdir(face_dir):
|
face_dict[name] = []
|
||||||
|
|
||||||
|
for file in sorted(
|
||||||
|
os.listdir(face_dir),
|
||||||
|
key=lambda f: os.path.getctime(os.path.join(face_dir, f)),
|
||||||
|
reverse=True,
|
||||||
|
):
|
||||||
face_dict[name].append(file)
|
face_dict[name].append(file)
|
||||||
|
|
||||||
return JSONResponse(status_code=200, content=face_dict)
|
return JSONResponse(status_code=200, content=face_dict)
|
||||||
@ -81,6 +85,10 @@ def train_face(request: Request, name: str, body: dict = None):
|
|||||||
new_name = f"{name}-{rand_id}.webp"
|
new_name = f"{name}-{rand_id}.webp"
|
||||||
new_file = os.path.join(FACE_DIR, f"{name}/{new_name}")
|
new_file = os.path.join(FACE_DIR, f"{name}/{new_name}")
|
||||||
shutil.move(training_file, new_file)
|
shutil.move(training_file, new_file)
|
||||||
|
|
||||||
|
context: EmbeddingsContext = request.app.embeddings
|
||||||
|
context.clear_face_classifier()
|
||||||
|
|
||||||
return JSONResponse(
|
return JSONResponse(
|
||||||
content=(
|
content=(
|
||||||
{
|
{
|
||||||
|
@ -9,6 +9,7 @@ SOCKET_REP_REQ = "ipc:///tmp/cache/embeddings"
|
|||||||
|
|
||||||
|
|
||||||
class EmbeddingsRequestEnum(Enum):
|
class EmbeddingsRequestEnum(Enum):
|
||||||
|
clear_face_classifier = "clear_face_classifier"
|
||||||
embed_description = "embed_description"
|
embed_description = "embed_description"
|
||||||
embed_thumbnail = "embed_thumbnail"
|
embed_thumbnail = "embed_thumbnail"
|
||||||
generate_search = "generate_search"
|
generate_search = "generate_search"
|
||||||
|
@ -32,9 +32,12 @@ class RealTimeProcessorApi(ABC):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def handle_request(self, request_data: dict[str, any]) -> dict[str, any] | None:
|
def handle_request(
|
||||||
|
self, topic: str, request_data: dict[str, any]
|
||||||
|
) -> dict[str, any] | None:
|
||||||
"""Handle metadata requests.
|
"""Handle metadata requests.
|
||||||
Args:
|
Args:
|
||||||
|
topic (str): topic that dictates what work is requested.
|
||||||
request_data (dict): containing data about requested change to process.
|
request_data (dict): containing data about requested change to process.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
|
@ -146,7 +146,7 @@ class BirdProcessor(RealTimeProcessorApi):
|
|||||||
if resp.status_code == 200:
|
if resp.status_code == 200:
|
||||||
self.detected_birds[obj_data["id"]] = score
|
self.detected_birds[obj_data["id"]] = score
|
||||||
|
|
||||||
def handle_request(self, request_data):
|
def handle_request(self, topic, request_data):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def expire_object(self, object_id):
|
def expire_object(self, object_id):
|
||||||
|
@ -12,6 +12,7 @@ import cv2
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import requests
|
import requests
|
||||||
|
|
||||||
|
from frigate.comms.embeddings_updater import EmbeddingsRequestEnum
|
||||||
from frigate.config import FrigateConfig
|
from frigate.config import FrigateConfig
|
||||||
from frigate.const import FACE_DIR, FRIGATE_LOCALHOST, MODEL_CACHE_DIR
|
from frigate.const import FACE_DIR, FRIGATE_LOCALHOST, MODEL_CACHE_DIR
|
||||||
from frigate.util.image import area
|
from frigate.util.image import area
|
||||||
@ -353,45 +354,52 @@ class FaceProcessor(RealTimeProcessorApi):
|
|||||||
|
|
||||||
self.__update_metrics(datetime.datetime.now().timestamp() - start)
|
self.__update_metrics(datetime.datetime.now().timestamp() - start)
|
||||||
|
|
||||||
def handle_request(self, request_data) -> dict[str, any] | None:
|
def handle_request(self, topic, request_data) -> dict[str, any] | None:
|
||||||
rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
|
if topic == EmbeddingsRequestEnum.clear_face_classifier.value:
|
||||||
label = request_data["face_name"]
|
self.__clear_classifier()
|
||||||
id = f"{label}-{rand_id}"
|
elif topic == EmbeddingsRequestEnum.register_face.value:
|
||||||
|
rand_id = "".join(
|
||||||
if request_data.get("cropped"):
|
random.choices(string.ascii_lowercase + string.digits, k=6)
|
||||||
thumbnail = request_data["image"]
|
|
||||||
else:
|
|
||||||
img = cv2.imdecode(
|
|
||||||
np.frombuffer(base64.b64decode(request_data["image"]), dtype=np.uint8),
|
|
||||||
cv2.IMREAD_COLOR,
|
|
||||||
)
|
)
|
||||||
face_box = self.__detect_face(img)
|
label = request_data["face_name"]
|
||||||
|
id = f"{label}-{rand_id}"
|
||||||
|
|
||||||
if not face_box:
|
if request_data.get("cropped"):
|
||||||
return {
|
thumbnail = request_data["image"]
|
||||||
"message": "No face was detected.",
|
else:
|
||||||
"success": False,
|
img = cv2.imdecode(
|
||||||
}
|
np.frombuffer(
|
||||||
|
base64.b64decode(request_data["image"]), dtype=np.uint8
|
||||||
|
),
|
||||||
|
cv2.IMREAD_COLOR,
|
||||||
|
)
|
||||||
|
face_box = self.__detect_face(img)
|
||||||
|
|
||||||
face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]]
|
if not face_box:
|
||||||
ret, thumbnail = cv2.imencode(
|
return {
|
||||||
".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100]
|
"message": "No face was detected.",
|
||||||
)
|
"success": False,
|
||||||
|
}
|
||||||
|
|
||||||
# write face to library
|
face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]]
|
||||||
folder = os.path.join(FACE_DIR, label)
|
_, thumbnail = cv2.imencode(
|
||||||
file = os.path.join(folder, f"{id}.webp")
|
".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100]
|
||||||
os.makedirs(folder, exist_ok=True)
|
)
|
||||||
|
|
||||||
# save face image
|
# write face to library
|
||||||
with open(file, "wb") as output:
|
folder = os.path.join(FACE_DIR, label)
|
||||||
output.write(thumbnail.tobytes())
|
file = os.path.join(folder, f"{id}.webp")
|
||||||
|
os.makedirs(folder, exist_ok=True)
|
||||||
|
|
||||||
self.__clear_classifier()
|
# save face image
|
||||||
return {
|
with open(file, "wb") as output:
|
||||||
"message": "Successfully registered face.",
|
output.write(thumbnail.tobytes())
|
||||||
"success": True,
|
|
||||||
}
|
self.__clear_classifier()
|
||||||
|
return {
|
||||||
|
"message": "Successfully registered face.",
|
||||||
|
"success": True,
|
||||||
|
}
|
||||||
|
|
||||||
def expire_object(self, object_id: str):
|
def expire_object(self, object_id: str):
|
||||||
if object_id in self.detected_faces:
|
if object_id in self.detected_faces:
|
||||||
|
@ -211,6 +211,11 @@ class EmbeddingsContext:
|
|||||||
|
|
||||||
return self.db.execute_sql(sql_query).fetchall()
|
return self.db.execute_sql(sql_query).fetchall()
|
||||||
|
|
||||||
|
def clear_face_classifier(self) -> None:
|
||||||
|
self.requestor.send_data(
|
||||||
|
EmbeddingsRequestEnum.clear_face_classifier.value, None
|
||||||
|
)
|
||||||
|
|
||||||
def delete_face_ids(self, face: str, ids: list[str]) -> None:
|
def delete_face_ids(self, face: str, ids: list[str]) -> None:
|
||||||
folder = os.path.join(FACE_DIR, face)
|
folder = os.path.join(FACE_DIR, face)
|
||||||
for id in ids:
|
for id in ids:
|
||||||
|
@ -140,7 +140,7 @@ class EmbeddingMaintainer(threading.Thread):
|
|||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
for processor in self.processors:
|
for processor in self.processors:
|
||||||
resp = processor.handle_request(data)
|
resp = processor.handle_request(topic, data)
|
||||||
|
|
||||||
if resp is not None:
|
if resp is not None:
|
||||||
return resp
|
return resp
|
||||||
|
Loading…
Reference in New Issue
Block a user