mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-02-05 00:15:51 +01:00
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:
parent
f4501a2094
commit
bb51a21bed
@ -34,6 +34,7 @@ from frigate.const import (
|
||||
CLIPS_DIR,
|
||||
CONFIG_DIR,
|
||||
EXPORT_DIR,
|
||||
FACE_DIR,
|
||||
MODEL_CACHE_DIR,
|
||||
RECORD_DIR,
|
||||
SHM_FRAMES_VAR,
|
||||
@ -96,14 +97,19 @@ class FrigateApp:
|
||||
self.config = config
|
||||
|
||||
def ensure_dirs(self) -> None:
|
||||
for d in [
|
||||
dirs = [
|
||||
CONFIG_DIR,
|
||||
RECORD_DIR,
|
||||
f"{CLIPS_DIR}/cache",
|
||||
CACHE_DIR,
|
||||
MODEL_CACHE_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):
|
||||
logger.info(f"Creating directory: {d}")
|
||||
os.makedirs(d)
|
||||
|
@ -123,19 +123,6 @@ class Embeddings:
|
||||
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_classification_model = None
|
||||
self.lpr_recognition_model = None
|
||||
|
@ -100,19 +100,6 @@ class EmbeddingMaintainer(threading.Thread):
|
||||
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:
|
||||
"""Maintain a SQLite-vec database for semantic search."""
|
||||
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]:
|
||||
"""Detect faces in input image."""
|
||||
self.face_detector.setInputSize((input.shape[1], input.shape[0]))
|
||||
faces = self.face_detector.detect(input)
|
||||
faces = self.face_classifier.detect_faces(input)
|
||||
|
||||
if faces[1] is None:
|
||||
if faces is None or faces[1] is None:
|
||||
return None
|
||||
|
||||
face = None
|
||||
|
@ -51,12 +51,14 @@ class ModelDownloader:
|
||||
download_path: str,
|
||||
file_names: List[str],
|
||||
download_func: Callable[[str], None],
|
||||
complete_func: Callable[[], None] | None = None,
|
||||
silent: bool = False,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.download_path = download_path
|
||||
self.file_names = file_names
|
||||
self.download_func = download_func
|
||||
self.complete_func = complete_func
|
||||
self.silent = silent
|
||||
self.requestor = InterProcessRequestor()
|
||||
self.download_thread = None
|
||||
@ -97,6 +99,9 @@ class ModelDownloader:
|
||||
},
|
||||
)
|
||||
|
||||
if self.complete_func:
|
||||
self.complete_func()
|
||||
|
||||
self.requestor.stop()
|
||||
self.download_complete.set()
|
||||
|
||||
|
@ -2,7 +2,7 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@ -10,6 +10,7 @@ import onnxruntime as ort
|
||||
from playhouse.sqliteq import SqliteQueueDatabase
|
||||
|
||||
from frigate.config.semantic_search import FaceRecognitionConfig
|
||||
from frigate.const import MODEL_CACHE_DIR
|
||||
|
||||
try:
|
||||
import openvino as ov
|
||||
@ -162,22 +163,62 @@ class FaceClassificationModel:
|
||||
def __init__(self, config: FaceRecognitionConfig, db: SqliteQueueDatabase):
|
||||
self.config = config
|
||||
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"):
|
||||
self.landmark_detector.loadModel(
|
||||
"/config/model_cache/facedet/landmarkdet.yaml"
|
||||
)
|
||||
download_path = os.path.join(MODEL_CACHE_DIR, "facedet")
|
||||
self.model_files = {
|
||||
"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 = (
|
||||
cv2.face.LBPHFaceRecognizer_create(
|
||||
radius=2, threshold=(1 - config.min_score) * 1000
|
||||
if not all(
|
||||
os.path.exists(os.path.join(download_path, n))
|
||||
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.__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:
|
||||
if not self.landmark_detector:
|
||||
return None
|
||||
|
||||
labels = []
|
||||
faces = []
|
||||
|
||||
@ -203,6 +244,11 @@ class FaceClassificationModel:
|
||||
faces.append(img)
|
||||
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))
|
||||
|
||||
def __align_face(
|
||||
@ -267,7 +313,17 @@ class FaceClassificationModel:
|
||||
self.labeler = None
|
||||
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:
|
||||
self.__build_classifier()
|
||||
|
||||
|
@ -1,5 +1,6 @@
|
||||
import { baseUrl } from "@/api/baseUrl";
|
||||
import AddFaceIcon from "@/components/icons/AddFaceIcon";
|
||||
import ActivityIndicator from "@/components/indicators/activity-indicator";
|
||||
import UploadImageDialog from "@/components/overlay/dialog/UploadImageDialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
@ -18,6 +19,8 @@ import {
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
import useOptimisticState from "@/hooks/use-optimistic-state";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { FrigateConfig } from "@/types/frigateConfig";
|
||||
import axios from "axios";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { LuImagePlus, LuTrash2 } from "react-icons/lu";
|
||||
@ -25,6 +28,8 @@ import { toast } from "sonner";
|
||||
import useSWR from "swr";
|
||||
|
||||
export default function FaceLibrary() {
|
||||
const { data: config } = useSWR<FrigateConfig>("config");
|
||||
|
||||
// title
|
||||
|
||||
useEffect(() => {
|
||||
@ -108,6 +113,10 @@ export default function FaceLibrary() {
|
||||
[pageToggle, refreshFaces],
|
||||
);
|
||||
|
||||
if (!config) {
|
||||
return <ActivityIndicator />;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex size-full flex-col p-2">
|
||||
<Toaster />
|
||||
@ -120,7 +129,7 @@ export default function FaceLibrary() {
|
||||
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">
|
||||
<div ref={tabsRef} className="flex flex-row">
|
||||
<ToggleGroup
|
||||
@ -156,17 +165,24 @@ export default function FaceLibrary() {
|
||||
data-nav-item={item}
|
||||
aria-label={`Select ${item}`}
|
||||
>
|
||||
<div className="capitalize">{item}</div>
|
||||
<div className="capitalize">
|
||||
{item} ({faceData[item].length})
|
||||
</div>
|
||||
</ToggleGroupItem>
|
||||
))}
|
||||
</ToggleGroup>
|
||||
<ScrollBar orientation="horizontal" className="h-0" />
|
||||
</div>
|
||||
</ScrollArea>
|
||||
<Button className="flex gap-2" onClick={() => setUpload(true)}>
|
||||
<LuImagePlus className="size-7 rounded-md p-1 text-secondary-foreground" />
|
||||
Upload Image
|
||||
</Button>
|
||||
</div>
|
||||
{pageToggle &&
|
||||
(pageToggle == "train" ? (
|
||||
<TrainingGrid
|
||||
config={config}
|
||||
attemptImages={trainImages}
|
||||
faceNames={faces}
|
||||
onRefresh={refreshFaces}
|
||||
@ -175,7 +191,6 @@ export default function FaceLibrary() {
|
||||
<FaceGrid
|
||||
faceImages={faceImages}
|
||||
pageToggle={pageToggle}
|
||||
setUpload={setUpload}
|
||||
onRefresh={refreshFaces}
|
||||
/>
|
||||
))}
|
||||
@ -184,11 +199,13 @@ export default function FaceLibrary() {
|
||||
}
|
||||
|
||||
type TrainingGridProps = {
|
||||
config: FrigateConfig;
|
||||
attemptImages: string[];
|
||||
faceNames: string[];
|
||||
onRefresh: () => void;
|
||||
};
|
||||
function TrainingGrid({
|
||||
config,
|
||||
attemptImages,
|
||||
faceNames,
|
||||
onRefresh,
|
||||
@ -200,6 +217,7 @@ function TrainingGrid({
|
||||
key={image}
|
||||
image={image}
|
||||
faceNames={faceNames}
|
||||
threshold={config.face_recognition.threshold}
|
||||
onRefresh={onRefresh}
|
||||
/>
|
||||
))}
|
||||
@ -210,9 +228,15 @@ function TrainingGrid({
|
||||
type FaceAttemptProps = {
|
||||
image: string;
|
||||
faceNames: string[];
|
||||
threshold: number;
|
||||
onRefresh: () => void;
|
||||
};
|
||||
function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) {
|
||||
function FaceAttempt({
|
||||
image,
|
||||
faceNames,
|
||||
threshold,
|
||||
onRefresh,
|
||||
}: FaceAttemptProps) {
|
||||
const data = useMemo(() => {
|
||||
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 flex-col items-start text-xs text-primary-variant">
|
||||
<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 className="flex flex-row items-start justify-end gap-5 md:gap-4">
|
||||
<Tooltip>
|
||||
@ -327,15 +359,9 @@ function FaceAttempt({ image, faceNames, onRefresh }: FaceAttemptProps) {
|
||||
type FaceGridProps = {
|
||||
faceImages: string[];
|
||||
pageToggle: string;
|
||||
setUpload: (upload: boolean) => void;
|
||||
onRefresh: () => void;
|
||||
};
|
||||
function FaceGrid({
|
||||
faceImages,
|
||||
pageToggle,
|
||||
setUpload,
|
||||
onRefresh,
|
||||
}: FaceGridProps) {
|
||||
function FaceGrid({ faceImages, pageToggle, onRefresh }: FaceGridProps) {
|
||||
return (
|
||||
<div className="scrollbar-container flex flex-wrap gap-2 overflow-y-scroll">
|
||||
{faceImages.map((image: string) => (
|
||||
@ -346,9 +372,6 @@ function FaceGrid({
|
||||
onRefresh={onRefresh}
|
||||
/>
|
||||
))}
|
||||
<Button key="upload" className="size-40" onClick={() => setUpload(true)}>
|
||||
<LuImagePlus className="size-10" />
|
||||
</Button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
@ -290,6 +290,7 @@ export interface FrigateConfig {
|
||||
|
||||
face_recognition: {
|
||||
enabled: boolean;
|
||||
threshold: number;
|
||||
};
|
||||
|
||||
ffmpeg: {
|
||||
|
Loading…
Reference in New Issue
Block a user