From e75d239ff25481038ca537693b9796326ca17e72 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Thu, 2 Jan 2025 16:44:25 -0600 Subject: [PATCH] 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 --- frigate/api/classification.py | 44 ++++- frigate/embeddings/maintainer.py | 2 +- frigate/util/model.py | 19 +- web/src/components/icons/AddFaceIcon.tsx | 25 +++ web/src/pages/FaceLibrary.tsx | 213 ++++++++++++++++------- 5 files changed, 231 insertions(+), 72 deletions(-) create mode 100644 web/src/components/icons/AddFaceIcon.tsx diff --git a/frigate/api/classification.py b/frigate/api/classification.py index fe54bebe9..6405516e0 100644 --- a/frigate/api/classification.py +++ b/frigate/api/classification.py @@ -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 {} diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py index b3c9610e4..b5c050325 100644 --- a/frigate/embeddings/maintainer.py +++ b/frigate/embeddings/maintainer.py @@ -517,7 +517,7 @@ class EmbeddingMaintainer(threading.Thread): if self.config.face_recognition.save_attempts: # write face to library - folder = os.path.join(FACE_DIR, "debug") + folder = os.path.join(FACE_DIR, "train") file = os.path.join(folder, f"{id}-{sub_label}-{score}-{face_score}.webp") os.makedirs(folder, exist_ok=True) cv2.imwrite(file, face_frame) diff --git a/frigate/util/model.py b/frigate/util/model.py index c95bc48b9..19c4b4e76 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -163,7 +163,12 @@ class FaceClassificationModel: self.config = config self.db = db self.landmark_detector = cv2.face.createFacemarkLBF() - self.landmark_detector.loadModel("/config/model_cache/facedet/landmarkdet.yaml") + + if os.path.isfile("/config/model_cache/facedet/landmarkdet.yaml"): + self.landmark_detector.loadModel( + "/config/model_cache/facedet/landmarkdet.yaml" + ) + self.recognizer: cv2.face.LBPHFaceRecognizer = ( cv2.face.LBPHFaceRecognizer_create( radius=2, threshold=(1 - config.min_score) * 1000 @@ -178,13 +183,21 @@ class FaceClassificationModel: dir = "/media/frigate/clips/faces" for idx, name in enumerate(os.listdir(dir)): - if name == "debug": + if name == "train": + continue + + face_folder = os.path.join(dir, name) + + if not os.path.isdir(face_folder): continue self.label_map[idx] = name - face_folder = os.path.join(dir, name) for image in os.listdir(face_folder): img = cv2.imread(os.path.join(face_folder, image)) + + if img is None: + continue + img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img = self.__align_face(img, img.shape[1], img.shape[0]) faces.append(img) diff --git a/web/src/components/icons/AddFaceIcon.tsx b/web/src/components/icons/AddFaceIcon.tsx new file mode 100644 index 000000000..ce06120cc --- /dev/null +++ b/web/src/components/icons/AddFaceIcon.tsx @@ -0,0 +1,25 @@ +import { forwardRef } from "react"; +import { LuPlus, LuScanFace } from "react-icons/lu"; +import { cn } from "@/lib/utils"; + +type AddFaceIconProps = { + className?: string; + onClick?: () => void; +}; + +const AddFaceIcon = forwardRef( + ({ className, onClick }, ref) => { + return ( +
+ + +
+ ); + }, +); + +export default AddFaceIcon; diff --git a/web/src/pages/FaceLibrary.tsx b/web/src/pages/FaceLibrary.tsx index 6955fc2c9..de1610518 100644 --- a/web/src/pages/FaceLibrary.tsx +++ b/web/src/pages/FaceLibrary.tsx @@ -1,19 +1,36 @@ import { baseUrl } from "@/api/baseUrl"; -import Chip from "@/components/indicators/Chip"; +import AddFaceIcon from "@/components/icons/AddFaceIcon"; import UploadImageDialog from "@/components/overlay/dialog/UploadImageDialog"; import { Button } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuLabel, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; import { ScrollArea, ScrollBar } from "@/components/ui/scroll-area"; import { Toaster } from "@/components/ui/sonner"; import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; import useOptimisticState from "@/hooks/use-optimistic-state"; import axios from "axios"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { isDesktop } from "react-device-detect"; -import { LuImagePlus, LuTrash } from "react-icons/lu"; +import { LuImagePlus, LuTrash2 } from "react-icons/lu"; import { toast } from "sonner"; import useSWR from "swr"; export default function FaceLibrary() { + // title + + useEffect(() => { + document.title = "Face Library - Frigate"; + }, []); + const [page, setPage] = useState(); const [pageToggle, setPageToggle] = useOptimisticState(page, setPage, 100); const tabsRef = useRef(null); @@ -24,7 +41,7 @@ export default function FaceLibrary() { const faces = useMemo( () => - faceData ? Object.keys(faceData).filter((face) => face != "debug") : [], + faceData ? Object.keys(faceData).filter((face) => face != "train") : [], [faceData], ); const faceImages = useMemo( @@ -32,24 +49,24 @@ export default function FaceLibrary() { [pageToggle, faceData], ); - const faceAttempts = useMemo( - () => faceData?.["debug"] || [], + const trainImages = useMemo( + () => faceData?.["train"] || [], [faceData], ); useEffect(() => { if (!pageToggle) { - if (faceAttempts.length > 0) { - setPageToggle("attempts"); + if (trainImages.length > 0) { + setPageToggle("train"); } else if (faces) { setPageToggle(faces[0]); } - } else if (pageToggle == "attempts" && faceAttempts.length == 0) { + } else if (pageToggle == "train" && trainImages.length == 0) { setPageToggle(faces[0]); } // we need to listen on the value of the faces list // eslint-disable-next-line react-hooks/exhaustive-deps - }, [faceAttempts, faces]); + }, [trainImages, faces]); // upload @@ -117,15 +134,15 @@ export default function FaceLibrary() { } }} > - {faceAttempts.length > 0 && ( + {trainImages.length > 0 && ( <> -
Attempts
+
Train
|
@@ -148,8 +165,12 @@ export default function FaceLibrary() { {pageToggle && - (pageToggle == "attempts" ? ( - + (pageToggle == "train" ? ( + ) : ( void; }; -function AttemptsGrid({ attemptImages, onRefresh }: AttemptsGridProps) { +function TrainingGrid({ + attemptImages, + faceNames, + onRefresh, +}: TrainingGridProps) { return (
{attemptImages.map((image: string) => ( - + ))}
); @@ -178,11 +209,10 @@ function AttemptsGrid({ attemptImages, onRefresh }: AttemptsGridProps) { type FaceAttemptProps = { image: string; + faceNames: string[]; onRefresh: () => void; }; -function FaceAttempt({ image, onRefresh }: FaceAttemptProps) { - const [hovered, setHovered] = useState(false); - +function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) { const data = useMemo(() => { const parts = image.split("-"); @@ -193,9 +223,36 @@ function FaceAttempt({ image, onRefresh }: FaceAttemptProps) { }; }, [image]); + const onTrainAttempt = useCallback( + (trainName: string) => { + axios + .post(`/faces/train/${trainName}/classify`, { training_file: image }) + .then((resp) => { + if (resp.status == 200) { + toast.success(`Successfully trained face.`, { + position: "top-center", + }); + onRefresh(); + } + }) + .catch((error) => { + if (error.response?.data?.message) { + toast.error(`Failed to train: ${error.response.data.message}`, { + position: "top-center", + }); + } else { + toast.error(`Failed to train: ${error.message}`, { + position: "top-center", + }); + } + }); + }, + [image, onRefresh], + ); + const onDelete = useCallback(() => { axios - .post(`/faces/debug/delete`, { ids: [image] }) + .post(`/faces/train/delete`, { ids: [image] }) .then((resp) => { if (resp.status == 200) { toast.success(`Successfully deleted face.`, { @@ -218,28 +275,50 @@ function FaceAttempt({ image, onRefresh }: FaceAttemptProps) { }, [image, onRefresh]); return ( -
setHovered(true) : undefined} - onMouseLeave={isDesktop ? () => setHovered(false) : undefined} - onClick={isDesktop ? undefined : () => setHovered(!hovered)} - > - {hovered && ( -
- onDelete()} - > - - +
+
+ +
+
+
+
+
{data.name}
+
{Number.parseFloat(data.score) * 100}%
+
+
+ + + + + + + + + Train Face as: + {faceNames.map((faceName) => ( + onTrainAttempt(faceName)} + > + {faceName} + + ))} + + + Train Face as Person + + + + + + Delete Face Attempt + +
- )} -
- -
{`${data.name}: ${data.score}`}
); @@ -280,8 +359,6 @@ type FaceImageProps = { onRefresh: () => void; }; function FaceImage({ name, image, onRefresh }: FaceImageProps) { - const [hovered, setHovered] = useState(false); - const onDelete = useCallback(() => { axios .post(`/faces/${name}/delete`, { ids: [image] }) @@ -307,26 +384,28 @@ function FaceImage({ name, image, onRefresh }: FaceImageProps) { }, [name, image, onRefresh]); return ( -
setHovered(true) : undefined} - onMouseLeave={isDesktop ? () => setHovered(false) : undefined} - onClick={isDesktop ? undefined : () => setHovered(!hovered)} - > - {hovered && ( -
- onDelete()} - > - - +
+
+ +
+
+
+
+
{name}
+
+
+ + + + + Delete Face Attempt + +
- )} - +
); }