mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-03-16 02:17:46 +01:00
Implement face recognition training in UI (#15786)
* Rename debug to train * Add api to train image as person * Cleanup model running * Formatting * Fix * Set face recognition page title
This commit is contained in:
@@ -2,6 +2,9 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import string
|
||||
|
||||
from fastapi import APIRouter, Request, UploadFile
|
||||
from fastapi.responses import JSONResponse
|
||||
@@ -22,7 +25,13 @@ def get_faces():
|
||||
|
||||
for name in os.listdir(FACE_DIR):
|
||||
face_dict[name] = []
|
||||
for file in os.listdir(os.path.join(FACE_DIR, name)):
|
||||
|
||||
face_dir = os.path.join(FACE_DIR, name)
|
||||
|
||||
if not os.path.isdir(face_dir):
|
||||
continue
|
||||
|
||||
for file in os.listdir(face_dir):
|
||||
face_dict[name].append(file)
|
||||
|
||||
return JSONResponse(status_code=200, content=face_dict)
|
||||
@@ -38,6 +47,39 @@ async def register_face(request: Request, name: str, file: UploadFile):
|
||||
)
|
||||
|
||||
|
||||
@router.post("/faces/train/{name}/classify")
|
||||
def train_face(name: str, body: dict = None):
|
||||
json: dict[str, any] = body or {}
|
||||
training_file = os.path.join(
|
||||
FACE_DIR, f"train/{sanitize_filename(json.get('training_file', ''))}"
|
||||
)
|
||||
|
||||
if not training_file or not os.path.isfile(training_file):
|
||||
return JSONResponse(
|
||||
content=(
|
||||
{
|
||||
"success": False,
|
||||
"message": f"Invalid filename or no file exists: {training_file}",
|
||||
}
|
||||
),
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
|
||||
new_name = f"{name}-{rand_id}.webp"
|
||||
new_file = os.path.join(FACE_DIR, f"{name}/{new_name}")
|
||||
shutil.move(training_file, new_file)
|
||||
return JSONResponse(
|
||||
content=(
|
||||
{
|
||||
"success": True,
|
||||
"message": f"Successfully saved {training_file} as {new_name}.",
|
||||
}
|
||||
),
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/faces/{name}/delete")
|
||||
def deregister_faces(request: Request, name: str, body: dict = None):
|
||||
json: dict[str, any] = body or {}
|
||||
|
||||
Reference in New Issue
Block a user