Fix facedet download (#15811)

* Support downloading face models

* Handle download and loading correctly

* Add face dir creation

* Fix error

* Fix

* Formatting

* Move upload to button

* Show number of faces in library for each name

* Add text color for score

* Cleanup
This commit is contained in:
Nicolas Mowen 2025-01-04 15:21:47 -06:00
parent f4501a2094
commit bb51a21bed
7 changed files with 121 additions and 57 deletions

View File

@ -34,6 +34,7 @@ from frigate.const import (
CLIPS_DIR, CLIPS_DIR,
CONFIG_DIR, CONFIG_DIR,
EXPORT_DIR, EXPORT_DIR,
FACE_DIR,
MODEL_CACHE_DIR, MODEL_CACHE_DIR,
RECORD_DIR, RECORD_DIR,
SHM_FRAMES_VAR, SHM_FRAMES_VAR,
@ -96,14 +97,19 @@ class FrigateApp:
self.config = config self.config = config
def ensure_dirs(self) -> None: def ensure_dirs(self) -> None:
for d in [ dirs = [
CONFIG_DIR, CONFIG_DIR,
RECORD_DIR, RECORD_DIR,
f"{CLIPS_DIR}/cache", f"{CLIPS_DIR}/cache",
CACHE_DIR, CACHE_DIR,
MODEL_CACHE_DIR, MODEL_CACHE_DIR,
EXPORT_DIR, EXPORT_DIR,
]: ]
if self.config.face_recognition.enabled:
dirs.append(FACE_DIR)
for d in dirs:
if not os.path.exists(d) and not os.path.islink(d): if not os.path.exists(d) and not os.path.islink(d):
logger.info(f"Creating directory: {d}") logger.info(f"Creating directory: {d}")
os.makedirs(d) os.makedirs(d)

View File

@ -123,19 +123,6 @@ class Embeddings:
device="GPU" if config.semantic_search.model_size == "large" else "CPU", device="GPU" if config.semantic_search.model_size == "large" else "CPU",
) )
if self.config.face_recognition.enabled:
self.face_embedding = GenericONNXEmbedding(
model_name="facedet",
model_file="facedet.onnx",
download_urls={
"facedet.onnx": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/facedet.onnx",
"landmarkdet.yaml": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/landmarkdet.yaml",
},
model_size="small",
model_type=ModelTypeEnum.face,
requestor=self.requestor,
)
self.lpr_detection_model = None self.lpr_detection_model = None
self.lpr_classification_model = None self.lpr_classification_model = None
self.lpr_recognition_model = None self.lpr_recognition_model = None

View File

@ -100,19 +100,6 @@ class EmbeddingMaintainer(threading.Thread):
self.lpr_config, self.requestor, self.embeddings self.lpr_config, self.requestor, self.embeddings
) )
@property
def face_detector(self) -> cv2.FaceDetectorYN:
# Lazily create the classifier.
if "face_detector" not in self.__dict__:
self.__dict__["face_detector"] = cv2.FaceDetectorYN.create(
"/config/model_cache/facedet/facedet.onnx",
config="",
input_size=(320, 320),
score_threshold=0.8,
nms_threshold=0.3,
)
return self.__dict__["face_detector"]
def run(self) -> None: def run(self) -> None:
"""Maintain a SQLite-vec database for semantic search.""" """Maintain a SQLite-vec database for semantic search."""
while not self.stop_event.is_set(): while not self.stop_event.is_set():
@ -395,10 +382,9 @@ class EmbeddingMaintainer(threading.Thread):
def _detect_face(self, input: np.ndarray) -> tuple[int, int, int, int]: def _detect_face(self, input: np.ndarray) -> tuple[int, int, int, int]:
"""Detect faces in input image.""" """Detect faces in input image."""
self.face_detector.setInputSize((input.shape[1], input.shape[0])) faces = self.face_classifier.detect_faces(input)
faces = self.face_detector.detect(input)
if faces[1] is None: if faces is None or faces[1] is None:
return None return None
face = None face = None

View File

@ -51,12 +51,14 @@ class ModelDownloader:
download_path: str, download_path: str,
file_names: List[str], file_names: List[str],
download_func: Callable[[str], None], download_func: Callable[[str], None],
complete_func: Callable[[], None] | None = None,
silent: bool = False, silent: bool = False,
): ):
self.model_name = model_name self.model_name = model_name
self.download_path = download_path self.download_path = download_path
self.file_names = file_names self.file_names = file_names
self.download_func = download_func self.download_func = download_func
self.complete_func = complete_func
self.silent = silent self.silent = silent
self.requestor = InterProcessRequestor() self.requestor = InterProcessRequestor()
self.download_thread = None self.download_thread = None
@ -97,6 +99,9 @@ class ModelDownloader:
}, },
) )
if self.complete_func:
self.complete_func()
self.requestor.stop() self.requestor.stop()
self.download_complete.set() self.download_complete.set()

View File

@ -2,7 +2,7 @@
import logging import logging
import os import os
from typing import Any, Optional from typing import Any
import cv2 import cv2
import numpy as np import numpy as np
@ -10,6 +10,7 @@ import onnxruntime as ort
from playhouse.sqliteq import SqliteQueueDatabase from playhouse.sqliteq import SqliteQueueDatabase
from frigate.config.semantic_search import FaceRecognitionConfig from frigate.config.semantic_search import FaceRecognitionConfig
from frigate.const import MODEL_CACHE_DIR
try: try:
import openvino as ov import openvino as ov
@ -162,22 +163,62 @@ class FaceClassificationModel:
def __init__(self, config: FaceRecognitionConfig, db: SqliteQueueDatabase): def __init__(self, config: FaceRecognitionConfig, db: SqliteQueueDatabase):
self.config = config self.config = config
self.db = db self.db = db
self.landmark_detector = cv2.face.createFacemarkLBF() self.face_detector: cv2.FaceDetectorYN = None
self.landmark_detector: cv2.face.FacemarkLBF = None
self.face_recognizer: cv2.face.LBPHFaceRecognizer = None
if os.path.isfile("/config/model_cache/facedet/landmarkdet.yaml"): download_path = os.path.join(MODEL_CACHE_DIR, "facedet")
self.landmark_detector.loadModel( self.model_files = {
"/config/model_cache/facedet/landmarkdet.yaml" "facedet.onnx": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/facedet.onnx",
) "landmarkdet.yaml": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/landmarkdet.yaml",
}
self.recognizer: cv2.face.LBPHFaceRecognizer = ( if not all(
cv2.face.LBPHFaceRecognizer_create( os.path.exists(os.path.join(download_path, n))
radius=2, threshold=(1 - config.min_score) * 1000 for n in self.model_files.keys()
) ):
# conditionally import ModelDownloader
from frigate.util.downloader import ModelDownloader
self.downloader = ModelDownloader(
model_name="facedet",
download_path=download_path,
file_names=self.model_files.keys(),
download_func=self.__download_models,
complete_func=self.__build_detector,
) )
self.downloader.ensure_model_files()
else:
self.__build_detector()
self.label_map: dict[int, str] = {} self.label_map: dict[int, str] = {}
self.__build_classifier() self.__build_classifier()
def __download_models(self, path: str) -> None:
try:
file_name = os.path.basename(path)
# conditionally import ModelDownloader
from frigate.util.downloader import ModelDownloader
ModelDownloader.download_from_url(self.model_files[file_name], path)
except Exception as e:
logger.error(f"Failed to download {path}: {e}")
def __build_detector(self) -> None:
self.face_detector = cv2.FaceDetectorYN.create(
"/config/model_cache/facedet/facedet.onnx",
config="",
input_size=(320, 320),
score_threshold=0.8,
nms_threshold=0.3,
)
self.landmark_detector = cv2.face.createFacemarkLBF()
self.landmark_detector.loadModel("/config/model_cache/facedet/landmarkdet.yaml")
def __build_classifier(self) -> None: def __build_classifier(self) -> None:
if not self.landmark_detector:
return None
labels = [] labels = []
faces = [] faces = []
@ -203,6 +244,11 @@ class FaceClassificationModel:
faces.append(img) faces.append(img)
labels.append(idx) labels.append(idx)
self.recognizer: cv2.face.LBPHFaceRecognizer = (
cv2.face.LBPHFaceRecognizer_create(
radius=2, threshold=(1 - self.config.min_score) * 1000
)
)
self.recognizer.train(faces, np.array(labels)) self.recognizer.train(faces, np.array(labels))
def __align_face( def __align_face(
@ -267,7 +313,17 @@ class FaceClassificationModel:
self.labeler = None self.labeler = None
self.label_map = {} self.label_map = {}
def classify_face(self, face_image: np.ndarray) -> Optional[tuple[str, float]]: def detect_faces(self, input: np.ndarray) -> tuple[int, cv2.typing.MatLike] | None:
if not self.face_detector:
return None
self.face_detector.setInputSize((input.shape[1], input.shape[0]))
return self.face_detector.detect(input)
def classify_face(self, face_image: np.ndarray) -> tuple[str, float] | None:
if not self.landmark_detector:
return None
if not self.label_map: if not self.label_map:
self.__build_classifier() self.__build_classifier()

View File

@ -1,5 +1,6 @@
import { baseUrl } from "@/api/baseUrl"; import { baseUrl } from "@/api/baseUrl";
import AddFaceIcon from "@/components/icons/AddFaceIcon"; import AddFaceIcon from "@/components/icons/AddFaceIcon";
import ActivityIndicator from "@/components/indicators/activity-indicator";
import UploadImageDialog from "@/components/overlay/dialog/UploadImageDialog"; import UploadImageDialog from "@/components/overlay/dialog/UploadImageDialog";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { import {
@ -18,6 +19,8 @@ import {
TooltipTrigger, TooltipTrigger,
} from "@/components/ui/tooltip"; } from "@/components/ui/tooltip";
import useOptimisticState from "@/hooks/use-optimistic-state"; import useOptimisticState from "@/hooks/use-optimistic-state";
import { cn } from "@/lib/utils";
import { FrigateConfig } from "@/types/frigateConfig";
import axios from "axios"; import axios from "axios";
import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { LuImagePlus, LuTrash2 } from "react-icons/lu"; import { LuImagePlus, LuTrash2 } from "react-icons/lu";
@ -25,6 +28,8 @@ import { toast } from "sonner";
import useSWR from "swr"; import useSWR from "swr";
export default function FaceLibrary() { export default function FaceLibrary() {
const { data: config } = useSWR<FrigateConfig>("config");
// title // title
useEffect(() => { useEffect(() => {
@ -108,6 +113,10 @@ export default function FaceLibrary() {
[pageToggle, refreshFaces], [pageToggle, refreshFaces],
); );
if (!config) {
return <ActivityIndicator />;
}
return ( return (
<div className="flex size-full flex-col p-2"> <div className="flex size-full flex-col p-2">
<Toaster /> <Toaster />
@ -120,7 +129,7 @@ export default function FaceLibrary() {
onSave={onUploadImage} onSave={onUploadImage}
/> />
<div className="relative flex h-11 w-full items-center justify-between"> <div className="relative mb-2 flex h-11 w-full items-center justify-between">
<ScrollArea className="w-full whitespace-nowrap"> <ScrollArea className="w-full whitespace-nowrap">
<div ref={tabsRef} className="flex flex-row"> <div ref={tabsRef} className="flex flex-row">
<ToggleGroup <ToggleGroup
@ -156,17 +165,24 @@ export default function FaceLibrary() {
data-nav-item={item} data-nav-item={item}
aria-label={`Select ${item}`} aria-label={`Select ${item}`}
> >
<div className="capitalize">{item}</div> <div className="capitalize">
{item} ({faceData[item].length})
</div>
</ToggleGroupItem> </ToggleGroupItem>
))} ))}
</ToggleGroup> </ToggleGroup>
<ScrollBar orientation="horizontal" className="h-0" /> <ScrollBar orientation="horizontal" className="h-0" />
</div> </div>
</ScrollArea> </ScrollArea>
<Button className="flex gap-2" onClick={() => setUpload(true)}>
<LuImagePlus className="size-7 rounded-md p-1 text-secondary-foreground" />
Upload Image
</Button>
</div> </div>
{pageToggle && {pageToggle &&
(pageToggle == "train" ? ( (pageToggle == "train" ? (
<TrainingGrid <TrainingGrid
config={config}
attemptImages={trainImages} attemptImages={trainImages}
faceNames={faces} faceNames={faces}
onRefresh={refreshFaces} onRefresh={refreshFaces}
@ -175,7 +191,6 @@ export default function FaceLibrary() {
<FaceGrid <FaceGrid
faceImages={faceImages} faceImages={faceImages}
pageToggle={pageToggle} pageToggle={pageToggle}
setUpload={setUpload}
onRefresh={refreshFaces} onRefresh={refreshFaces}
/> />
))} ))}
@ -184,11 +199,13 @@ export default function FaceLibrary() {
} }
type TrainingGridProps = { type TrainingGridProps = {
config: FrigateConfig;
attemptImages: string[]; attemptImages: string[];
faceNames: string[]; faceNames: string[];
onRefresh: () => void; onRefresh: () => void;
}; };
function TrainingGrid({ function TrainingGrid({
config,
attemptImages, attemptImages,
faceNames, faceNames,
onRefresh, onRefresh,
@ -200,6 +217,7 @@ function TrainingGrid({
key={image} key={image}
image={image} image={image}
faceNames={faceNames} faceNames={faceNames}
threshold={config.face_recognition.threshold}
onRefresh={onRefresh} onRefresh={onRefresh}
/> />
))} ))}
@ -210,9 +228,15 @@ function TrainingGrid({
type FaceAttemptProps = { type FaceAttemptProps = {
image: string; image: string;
faceNames: string[]; faceNames: string[];
threshold: number;
onRefresh: () => void; onRefresh: () => void;
}; };
function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) { function FaceAttempt({
image,
faceNames,
threshold,
onRefresh,
}: FaceAttemptProps) {
const data = useMemo(() => { const data = useMemo(() => {
const parts = image.split("-"); const parts = image.split("-");
@ -283,7 +307,15 @@ function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) {
<div className="flex w-full flex-row items-center justify-between gap-2"> <div className="flex w-full flex-row items-center justify-between gap-2">
<div className="flex flex-col items-start text-xs text-primary-variant"> <div className="flex flex-col items-start text-xs text-primary-variant">
<div className="capitalize">{data.name}</div> <div className="capitalize">{data.name}</div>
<div>{Number.parseFloat(data.score) * 100}%</div> <div
className={cn(
Number.parseFloat(data.score) > threshold
? "text-success"
: "text-danger",
)}
>
{Number.parseFloat(data.score) * 100}%
</div>
</div> </div>
<div className="flex flex-row items-start justify-end gap-5 md:gap-4"> <div className="flex flex-row items-start justify-end gap-5 md:gap-4">
<Tooltip> <Tooltip>
@ -327,15 +359,9 @@ function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) {
type FaceGridProps = { type FaceGridProps = {
faceImages: string[]; faceImages: string[];
pageToggle: string; pageToggle: string;
setUpload: (upload: boolean) => void;
onRefresh: () => void; onRefresh: () => void;
}; };
function FaceGrid({ function FaceGrid({ faceImages, pageToggle, onRefresh }: FaceGridProps) {
faceImages,
pageToggle,
setUpload,
onRefresh,
}: FaceGridProps) {
return ( return (
<div className="scrollbar-container flex flex-wrap gap-2 overflow-y-scroll"> <div className="scrollbar-container flex flex-wrap gap-2 overflow-y-scroll">
{faceImages.map((image: string) => ( {faceImages.map((image: string) => (
@ -346,9 +372,6 @@ function FaceGrid({
onRefresh={onRefresh} onRefresh={onRefresh}
/> />
))} ))}
<Button key="upload" className="size-40" onClick={() => setUpload(true)}>
<LuImagePlus className="size-10" />
</Button>
</div> </div>
); );
} }

View File

@ -290,6 +290,7 @@ export interface FrigateConfig {
face_recognition: { face_recognition: {
enabled: boolean; enabled: boolean;
threshold: number;
}; };
ffmpeg: { ffmpeg: {