From b8216d05362548d19e3d1c5ca3bad625f15812b4 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Mon, 10 Nov 2025 05:51:03 -0700
Subject: [PATCH] Handle case where classification images are deleted

---
 frigate/api/classification.py                        | 5 +++++
 web/public/locales/en/views/classificationModel.json | 1 +
 web/src/views/classification/ModelTrainingView.tsx   | 9 +++++----
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/frigate/api/classification.py b/frigate/api/classification.py
index 87de52884..674c726a6 100644
--- a/frigate/api/classification.py
+++ b/frigate/api/classification.py
@@ -595,9 +595,13 @@ def get_classification_dataset(name: str):
             "last_training_image_count": 0,
             "current_image_count": current_image_count,
             "new_images_count": current_image_count,
+            "dataset_changed": current_image_count > 0,
         }
     else:
         last_training_count = metadata.get("last_training_image_count", 0)
+        # Dataset has changed if count is different (either added or deleted images)
+        dataset_changed = current_image_count != last_training_count
+        # Only show positive count for new images (ignore deletions in the count display)
         new_images_count = max(0, current_image_count - last_training_count)
         training_metadata = {
             "has_trained": True,
@@ -605,6 +609,7 @@ def get_classification_dataset(name: str):
             "last_training_image_count": last_training_count,
             "current_image_count": current_image_count,
             "new_images_count": new_images_count,
+            "dataset_changed": dataset_changed,
         }
 
     return JSONResponse(
diff --git a/web/public/locales/en/views/classificationModel.json b/web/public/locales/en/views/classificationModel.json
index 2bae0c0ce..f77d8067a 100644
--- a/web/public/locales/en/views/classificationModel.json
+++ b/web/public/locales/en/views/classificationModel.json
@@ -16,6 +16,7 @@
   "tooltip": {
     "trainingInProgress": "Model is currently training",
     "noNewImages": "No new images to train. Classify more images in the dataset first.",
+    "noChanges": "No changes to the dataset since last training.",
     "modelNotReady": "Model is not ready for training"
   },
   "toast": {
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index 6a3e680f9..61c99b15f 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -126,6 +126,7 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
       last_training_image_count: number;
       current_image_count: number;
       new_images_count: number;
+      dataset_changed: boolean;
     } | null;
   }>(`classification/${model.name}/dataset`);
 
@@ -445,7 +446,7 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
                   variant={modelState == "failed" ? "destructive" : "select"}
                   disabled={
                     (modelState != "complete" && modelState != "failed") ||
-                    (trainingMetadata?.new_images_count ?? 0) === 0
+                    !trainingMetadata?.dataset_changed
                   }
                 >
                   {modelState == "training" ? (
@@ -466,14 +467,14 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
                   )}
                 </Button>
               </TooltipTrigger>
-              {((trainingMetadata?.new_images_count ?? 0) === 0 ||
+              {(!trainingMetadata?.dataset_changed ||
                 (modelState != "complete" && modelState != "failed")) && (
                 <TooltipPortal>
                   <TooltipContent>
                     {modelState == "training"
                       ? t("tooltip.trainingInProgress")
-                      : trainingMetadata?.new_images_count === 0
-                        ? t("tooltip.noNewImages")
+                      : !trainingMetadata?.dataset_changed
+                        ? t("tooltip.noChanges")
                         : t("tooltip.modelNotReady")}
                   </TooltipContent>
                 </TooltipPortal>