From ca91d3d82b2aca048e5552d3888fe55f8b78b661 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 6 Nov 2025 15:55:05 -0700
Subject: [PATCH] Fix deletion of classification images and library

---
 frigate/api/classification.py                 |  9 ++++---
 .../classification/ModelTrainingView.tsx      | 25 ++++++++++---------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/frigate/api/classification.py b/frigate/api/classification.py
index a167911c4..eee637585 100644
--- a/frigate/api/classification.py
+++ b/frigate/api/classification.py
@@ -662,8 +662,11 @@ def delete_classification_dataset_images(
         if os.path.isfile(file_path):
             os.unlink(file_path)
 
+    if os.path.exists(folder) and not os.listdir(folder):
+        os.rmdir(folder)
+
     return JSONResponse(
-        content=({"success": True, "message": "Successfully deleted faces."}),
+        content=({"success": True, "message": "Successfully deleted images."}),
         status_code=200,
     )
 
@@ -723,7 +726,7 @@ def categorize_classification_image(request: Request, name: str, body: dict = No
     os.unlink(training_file)
 
     return JSONResponse(
-        content=({"success": True, "message": "Successfully deleted faces."}),
+        content=({"success": True, "message": "Successfully categorized image."}),
         status_code=200,
     )
 
@@ -761,7 +764,7 @@ def delete_classification_train_images(request: Request, name: str, body: dict =
             os.unlink(file_path)
 
     return JSONResponse(
-        content=({"success": True, "message": "Successfully deleted faces."}),
+        content=({"success": True, "message": "Successfully deleted images."}),
         status_code=200,
     )
 
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index a27a06a9e..036fe0e56 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -183,11 +183,12 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
   );
 
   const onDelete = useCallback(
-    (ids: string[], isName: boolean = false) => {
+    (ids: string[], isName: boolean = false, category?: string) => {
+      const targetCategory = category || pageToggle;
       const api =
-        pageToggle == "train"
+        targetCategory == "train"
           ? `/classification/${model.name}/train/delete`
-          : `/classification/${model.name}/dataset/${pageToggle}/delete`;
+          : `/classification/${model.name}/dataset/${targetCategory}/delete`;
 
       axios
         .post(api, { ids })
@@ -432,7 +433,7 @@ type LibrarySelectorProps = {
   dataset: { [id: string]: string[] };
   trainImages: string[];
   setPageToggle: (toggle: string) => void;
-  onDelete: (ids: string[], isName: boolean) => void;
+  onDelete: (ids: string[], isName: boolean, category?: string) => void;
   onRename: (old_name: string, new_name: string) => void;
 };
 function LibrarySelector({
@@ -448,7 +449,7 @@ function LibrarySelector({
   // data
 
   const [confirmDelete, setConfirmDelete] = useState<string | null>(null);
-  const [renameClass, setRenameFace] = useState<string | null>(null);
+  const [renameClass, setRenameClass] = useState<string | null>(null);
   const pageTitle = useMemo(() => {
     if (pageToggle != "train") {
       return pageToggle;
@@ -463,12 +464,12 @@ function LibrarySelector({
 
   // interaction
 
-  const handleDeleteFace = useCallback(
+  const handleDeleteCategory = useCallback(
     (name: string) => {
-      // Get all image IDs for this face
+      // Get all image IDs for this category
       const imageIds = dataset?.[name] || [];
 
-      onDelete(imageIds, true);
+      onDelete(imageIds, true, name);
       setPageToggle("train");
     },
     [dataset, onDelete, setPageToggle],
@@ -476,7 +477,7 @@ function LibrarySelector({
 
   const handleSetOpen = useCallback(
     (open: boolean) => {
-      setRenameFace(open ? renameClass : null);
+      setRenameClass(open ? renameClass : null);
     },
     [renameClass],
   );
@@ -503,7 +504,7 @@ function LibrarySelector({
               className="text-white"
               onClick={() => {
                 if (confirmDelete) {
-                  handleDeleteFace(confirmDelete);
+                  handleDeleteCategory(confirmDelete);
                   setConfirmDelete(null);
                 }
               }}
@@ -521,7 +522,7 @@ function LibrarySelector({
         description={t("renameCategory.desc", { name: renameClass })}
         onSave={(newName) => {
           onRename(renameClass!, newName);
-          setRenameFace(null);
+          setRenameClass(null);
         }}
         defaultValue={renameClass || ""}
         regexPattern={/^[\p{L}\p{N}\s'_-]{1,50}$/u}
@@ -588,7 +589,7 @@ function LibrarySelector({
                       className="size-7 lg:opacity-0 lg:transition-opacity lg:group-hover:opacity-100"
                       onClick={(e) => {
                         e.stopPropagation();
-                        setRenameFace(id);
+                        setRenameClass(id);
                       }}
                     >
                       <LuPencil className="size-4 text-primary" />