Add api to train image as person

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
@@ -38,6 +41,38 @@ async def register_face(request: Request, name: str, file: UploadFile):
     )
 
 
+@router.post("/faces/{name}/train")
+def train_face(name: str, body: dict = None):
+    json: dict[str, any] = body or {}
+    file_name = sanitize_filename(json.get("training_file", ""))
+    training_file = os.path.join(FACE_DIR, f"train/{file_name}")
+
+    if not file_name or not os.path.isfile(training_file):
+        return JSONResponse(
+            content=(
+                {
+                    "success": False,
+                    "message": f"Invalid filename or no file exists: {file_name}",
+                }
+            ),
+            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 {file_name} 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 {}

web/src/pages/FaceLibrary.tsx

@@ -8,14 +8,13 @@ import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group";
 import {
   Tooltip,
   TooltipContent,
-  TooltipProvider,
   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, LuTrash, LuTrash2 } from "react-icons/lu";
 import { toast } from "sonner";
 import useSWR from "swr";
 
@@ -126,12 +125,12 @@ export default function FaceLibrary() {
               {trainImages.length > 0 && (
                 <>
                   <ToggleGroupItem
-                    value="attempts"
+                    value="train"
-                    className={`flex scroll-mx-10 items-center justify-between gap-2 ${pageToggle == "attempts" ? "" : "*:text-muted-foreground"}`}
+                    className={`flex scroll-mx-10 items-center justify-between gap-2 ${pageToggle == "train" ? "" : "*:text-muted-foreground"}`}
-                    data-nav-item="attempts"
+                    data-nav-item="train"
-                    aria-label="Select attempts"
+                    aria-label="Select train"
                   >
-                    <div>Attempts</div>
+                    <div>Train</div>
                   </ToggleGroupItem>
                   <div>|</div>
                 </>
@@ -154,7 +153,7 @@ export default function FaceLibrary() {
         </ScrollArea>
       </div>
       {pageToggle &&
-        (pageToggle == "attempts" ? (
+        (pageToggle == "train" ? (
           <TrainingGrid attemptImages={trainImages} onRefresh={refreshFaces} />
         ) : (
           <FaceGrid
@@ -187,8 +186,6 @@ type FaceAttemptProps = {
   onRefresh: () => void;
 };
 function FaceAttempt({ image, onRefresh }: FaceAttemptProps) {
-  const [hovered, setHovered] = useState(false);
-
   const data = useMemo(() => {
     const parts = image.split("-");
 
@@ -224,33 +221,28 @@ function FaceAttempt({ image, onRefresh }: FaceAttemptProps) {
   }, [image, onRefresh]);
 
   return (
-    <div
+    <div className="relative flex flex-col rounded-lg">
-      className="relative h-min"
+      <div className="w-full overflow-hidden rounded-t-lg border border-t-0 *:text-card-foreground">
-      onMouseEnter={isDesktop ? () => setHovered(true) : undefined}
+        <img className="h-40" src={`${baseUrl}clips/faces/train/${image}`} />
-      onMouseLeave={isDesktop ? () => setHovered(false) : undefined}
+      </div>
-      onClick={isDesktop ? undefined : () => setHovered(!hovered)}
+      <div className="rounded-b-lg bg-card p-2">
-    >
+        <div className="flex w-full flex-row items-center justify-between gap-2">
-      {hovered && (
+          <div className="flex flex-col items-start text-xs text-primary-variant">
-        <Tooltip>
+            <div className="capitalize">{data.name}</div>
-          <div className="absolute right-1 top-1">
+            <div>{Number.parseFloat(data.score) * 100}%</div>
-            <TooltipTrigger>
-              <Chip
-                className="cursor-pointer rounded-md bg-gray-500 bg-gradient-to-br from-gray-400 to-gray-500"
-                onClick={() => onDelete()}
-              >
-                <LuTrash className="size-4 fill-destructive text-destructive" />
-              </Chip>
-            </TooltipTrigger>
-            <TooltipContent>Delete Face Attempt</TooltipContent>
           </div>
-        </Tooltip>
+          <div className="flex flex-row items-start justify-end gap-5 md:gap-4">
-      )}
+            <Tooltip>
-      <div className="rounded-md bg-secondary">
+              <TooltipTrigger>
-        <img
+                <LuTrash2
-          className="h-40 rounded-md"
+                  className="size-5 cursor-pointer text-primary-variant hover:text-primary"
-          src={`${baseUrl}clips/faces/train/${image}`}
+                  onClick={onDelete}
-        />
+                />
-        <div className="p-2">{`${data.name}: ${data.score}`}</div>
+              </TooltipTrigger>
+              <TooltipContent>Delete Face Attempt</TooltipContent>
+            </Tooltip>
+          </div>
+        </div>
       </div>
     </div>
   );
@@ -291,8 +283,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] })
@@ -318,31 +308,28 @@ function FaceImage({ name, image, onRefresh }: FaceImageProps) {
   }, [name, image, onRefresh]);
 
   return (
-    <div
+    <div className="relative flex flex-col rounded-lg">
-      className="relative h-40"
+      <div className="w-full overflow-hidden rounded-t-lg border border-t-0 *:text-card-foreground">
-      onMouseEnter={isDesktop ? () => setHovered(true) : undefined}
+        <img className="h-40" src={`${baseUrl}clips/faces/${name}/${image}`} />
-      onMouseLeave={isDesktop ? () => setHovered(false) : undefined}
+      </div>
-      onClick={isDesktop ? undefined : () => setHovered(!hovered)}
+      <div className="rounded-b-lg bg-card p-2">
-    >
+        <div className="flex w-full flex-row items-center justify-between gap-2">
-      {hovered && (
+          <div className="flex flex-col items-start text-xs text-primary-variant">
-        <Tooltip>
+            <div className="capitalize">{name}</div>
-          <div className="absolute right-1 top-1">
-            <TooltipTrigger>
-              <Chip
-                className="cursor-pointer rounded-md bg-gray-500 bg-gradient-to-br from-gray-400 to-gray-500"
-                onClick={() => onDelete()}
-              >
-                <LuTrash className="size-4 fill-destructive text-destructive" />
-              </Chip>
-            </TooltipTrigger>
-            <TooltipContent>Delete Face</TooltipContent>
           </div>
-        </Tooltip>
+          <div className="flex flex-row items-start justify-end gap-5 md:gap-4">
-      )}
+            <Tooltip>
-      <img
+              <TooltipTrigger>
-        className="h-40 rounded-md"
+                <LuTrash2
-        src={`${baseUrl}clips/faces/${name}/${image}`}
+                  className="size-5 cursor-pointer text-primary-variant hover:text-primary"
-      />
+                  onClick={onDelete}
+                />
+              </TooltipTrigger>
+              <TooltipContent>Delete Face Attempt</TooltipContent>
+            </Tooltip>
+          </div>
+        </div>
+      </div>
     </div>
   );
 }