Improve classification (#20863)

frigate/util/classification.py

@@ -1,5 +1,7 @@
 """Util for classification models."""
 
+import datetime
+import json
 import logging
 import os
 import random
@@ -27,10 +29,96 @@ from frigate.util.process import FrigateProcess
 BATCH_SIZE = 16
 EPOCHS = 50
 LEARNING_RATE = 0.001
+TRAINING_METADATA_FILE = ".training_metadata.json"
 
 logger = logging.getLogger(__name__)
 
 
+def write_training_metadata(model_name: str, image_count: int) -> None:
+    """
+    Write training metadata to a hidden file in the model's clips directory.
+
+    Args:
+        model_name: Name of the classification model
+        image_count: Number of images used in training
+    """
+    clips_model_dir = os.path.join(CLIPS_DIR, model_name)
+    os.makedirs(clips_model_dir, exist_ok=True)
+
+    metadata_path = os.path.join(clips_model_dir, TRAINING_METADATA_FILE)
+    metadata = {
+        "last_training_date": datetime.datetime.now().isoformat(),
+        "last_training_image_count": image_count,
+    }
+
+    try:
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+        logger.info(f"Wrote training metadata for {model_name}: {image_count} images")
+    except Exception as e:
+        logger.error(f"Failed to write training metadata for {model_name}: {e}")
+
+
+def read_training_metadata(model_name: str) -> dict[str, any] | None:
+    """
+    Read training metadata from the hidden file in the model's clips directory.
+
+    Args:
+        model_name: Name of the classification model
+
+    Returns:
+        Dictionary with last_training_date and last_training_image_count, or None if not found
+    """
+    clips_model_dir = os.path.join(CLIPS_DIR, model_name)
+    metadata_path = os.path.join(clips_model_dir, TRAINING_METADATA_FILE)
+
+    if not os.path.exists(metadata_path):
+        return None
+
+    try:
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+        return metadata
+    except Exception as e:
+        logger.error(f"Failed to read training metadata for {model_name}: {e}")
+        return None
+
+
+def get_dataset_image_count(model_name: str) -> int:
+    """
+    Count the total number of images in the model's dataset directory.
+
+    Args:
+        model_name: Name of the classification model
+
+    Returns:
+        Total count of images across all categories
+    """
+    dataset_dir = os.path.join(CLIPS_DIR, model_name, "dataset")
+
+    if not os.path.exists(dataset_dir):
+        return 0
+
+    total_count = 0
+    try:
+        for category in os.listdir(dataset_dir):
+            category_dir = os.path.join(dataset_dir, category)
+            if not os.path.isdir(category_dir):
+                continue
+
+            image_files = [
+                f
+                for f in os.listdir(category_dir)
+                if f.lower().endswith((".webp", ".png", ".jpg", ".jpeg"))
+            ]
+            total_count += len(image_files)
+    except Exception as e:
+        logger.error(f"Failed to count dataset images for {model_name}: {e}")
+        return 0
+
+    return total_count
+
+
 class ClassificationTrainingProcess(FrigateProcess):
     def __init__(self, model_name: str) -> None:
         super().__init__(
@@ -42,7 +130,8 @@ class ClassificationTrainingProcess(FrigateProcess):
 
     def run(self) -> None:
         self.pre_run_setup()
-        self.__train_classification_model()
+        success = self.__train_classification_model()
+        exit(0 if success else 1)
 
     def __generate_representative_dataset_factory(self, dataset_dir: str):
         def generate_representative_dataset():
@@ -65,85 +154,117 @@ class ClassificationTrainingProcess(FrigateProcess):
     @redirect_output_to_logger(logger, logging.DEBUG)
     def __train_classification_model(self) -> bool:
         """Train a classification model."""
+        try:
+            # import in the function so that tensorflow is not initialized multiple times
+            import tensorflow as tf
+            from tensorflow.keras import layers, models, optimizers
+            from tensorflow.keras.applications import MobileNetV2
+            from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
-        # import in the function so that tensorflow is not initialized multiple times
-        import tensorflow as tf
-        from tensorflow.keras import layers, models, optimizers
-        from tensorflow.keras.applications import MobileNetV2
-        from tensorflow.keras.preprocessing.image import ImageDataGenerator
+            dataset_dir = os.path.join(CLIPS_DIR, self.model_name, "dataset")
+            model_dir = os.path.join(MODEL_CACHE_DIR, self.model_name)
+            os.makedirs(model_dir, exist_ok=True)
 
-        logger.info(f"Kicking off classification training for {self.model_name}.")
-        dataset_dir = os.path.join(CLIPS_DIR, self.model_name, "dataset")
-        model_dir = os.path.join(MODEL_CACHE_DIR, self.model_name)
-        os.makedirs(model_dir, exist_ok=True)
-        num_classes = len(
-            [
-                d
-                for d in os.listdir(dataset_dir)
-                if os.path.isdir(os.path.join(dataset_dir, d))
-            ]
-        )
+            num_classes = len(
+                [
+                    d
+                    for d in os.listdir(dataset_dir)
+                    if os.path.isdir(os.path.join(dataset_dir, d))
+                ]
+            )
 
-        # Start with imagenet base model with 35% of channels in each layer
-        base_model = MobileNetV2(
-            input_shape=(224, 224, 3),
-            include_top=False,
-            weights="imagenet",
-            alpha=0.35,
-        )
-        base_model.trainable = False  # Freeze pre-trained layers
+            if num_classes < 2:
+                logger.error(
+                    f"Training failed for {self.model_name}: Need at least 2 classes, found {num_classes}"
+                )
+                return False
 
-        model = models.Sequential(
-            [
-                base_model,
-                layers.GlobalAveragePooling2D(),
-                layers.Dense(128, activation="relu"),
-                layers.Dropout(0.3),
-                layers.Dense(num_classes, activation="softmax"),
-            ]
-        )
+            # Start with imagenet base model with 35% of channels in each layer
+            base_model = MobileNetV2(
+                input_shape=(224, 224, 3),
+                include_top=False,
+                weights="imagenet",
+                alpha=0.35,
+            )
+            base_model.trainable = False  # Freeze pre-trained layers
 
-        model.compile(
-            optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
-            loss="categorical_crossentropy",
-            metrics=["accuracy"],
-        )
+            model = models.Sequential(
+                [
+                    base_model,
+                    layers.GlobalAveragePooling2D(),
+                    layers.Dense(128, activation="relu"),
+                    layers.Dropout(0.3),
+                    layers.Dense(num_classes, activation="softmax"),
+                ]
+            )
 
-        # create training set
-        datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
-        train_gen = datagen.flow_from_directory(
-            dataset_dir,
-            target_size=(224, 224),
-            batch_size=BATCH_SIZE,
-            class_mode="categorical",
-            subset="training",
-        )
+            model.compile(
+                optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
+                loss="categorical_crossentropy",
+                metrics=["accuracy"],
+            )
 
-        # write labelmap
-        class_indices = train_gen.class_indices
-        index_to_class = {v: k for k, v in class_indices.items()}
-        sorted_classes = [index_to_class[i] for i in range(len(index_to_class))]
-        with open(os.path.join(model_dir, "labelmap.txt"), "w") as f:
-            for class_name in sorted_classes:
-                f.write(f"{class_name}\n")
+            # create training set
+            datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
+            train_gen = datagen.flow_from_directory(
+                dataset_dir,
+                target_size=(224, 224),
+                batch_size=BATCH_SIZE,
+                class_mode="categorical",
+                subset="training",
+            )
 
-        # train the model
-        model.fit(train_gen, epochs=EPOCHS, verbose=0)
+            total_images = train_gen.samples
+            logger.debug(
+                f"Training {self.model_name}: {total_images} images across {num_classes} classes"
+            )
 
-        # convert model to tflite
-        converter = tf.lite.TFLiteConverter.from_keras_model(model)
-        converter.optimizations = [tf.lite.Optimize.DEFAULT]
-        converter.representative_dataset = (
-            self.__generate_representative_dataset_factory(dataset_dir)
-        )
-        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-        converter.inference_input_type = tf.uint8
-        converter.inference_output_type = tf.uint8
-        tflite_model = converter.convert()
+            # write labelmap
+            class_indices = train_gen.class_indices
+            index_to_class = {v: k for k, v in class_indices.items()}
+            sorted_classes = [index_to_class[i] for i in range(len(index_to_class))]
+            with open(os.path.join(model_dir, "labelmap.txt"), "w") as f:
+                for class_name in sorted_classes:
+                    f.write(f"{class_name}\n")
 
-        # write model
-        with open(os.path.join(model_dir, "model.tflite"), "wb") as f:
-            f.write(tflite_model)
+            # train the model
+            logger.debug(f"Training {self.model_name} for {EPOCHS} epochs...")
+            model.fit(train_gen, epochs=EPOCHS, verbose=0)
+            logger.debug(f"Converting {self.model_name} to TFLite...")
+
+            # convert model to tflite
+            converter = tf.lite.TFLiteConverter.from_keras_model(model)
+            converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            converter.representative_dataset = (
+                self.__generate_representative_dataset_factory(dataset_dir)
+            )
+            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+            converter.inference_input_type = tf.uint8
+            converter.inference_output_type = tf.uint8
+            tflite_model = converter.convert()
+
+            # write model
+            model_path = os.path.join(model_dir, "model.tflite")
+            with open(model_path, "wb") as f:
+                f.write(tflite_model)
+
+            # verify model file was written successfully
+            if not os.path.exists(model_path) or os.path.getsize(model_path) == 0:
+                logger.error(
+                    f"Training failed for {self.model_name}: Model file was not created or is empty"
+                )
+                return False
+
+            # write training metadata with image count
+            dataset_image_count = get_dataset_image_count(self.model_name)
+            write_training_metadata(self.model_name, dataset_image_count)
+
+            logger.info(f"Finished training {self.model_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Training failed for {self.model_name}: {e}", exc_info=True)
+            return False
 
 
 def kickoff_model_training(
@@ -165,18 +286,36 @@ def kickoff_model_training(
     training_process.start()
     training_process.join()
 
-    # reload model and mark training as complete
-    embeddingRequestor.send_data(
-        EmbeddingsRequestEnum.reload_classification_model.value,
-        {"model_name": model_name},
-    )
-    requestor.send_data(
-        UPDATE_MODEL_STATE,
-        {
-            "model": model_name,
-            "state": ModelStatusTypesEnum.complete,
-        },
-    )
+    # check if training succeeded by examining the exit code
+    training_success = training_process.exitcode == 0
+
+    if training_success:
+        # reload model and mark training as complete
+        embeddingRequestor.send_data(
+            EmbeddingsRequestEnum.reload_classification_model.value,
+            {"model_name": model_name},
+        )
+        requestor.send_data(
+            UPDATE_MODEL_STATE,
+            {
+                "model": model_name,
+                "state": ModelStatusTypesEnum.complete,
+            },
+        )
+    else:
+        logger.error(
+            f"Training subprocess failed for {model_name} (exit code: {training_process.exitcode})"
+        )
+        # mark training as failed so UI shows error state
+        # don't reload the model since it failed
+        requestor.send_data(
+            UPDATE_MODEL_STATE,
+            {
+                "model": model_name,
+                "state": ModelStatusTypesEnum.failed,
+            },
+        )
+
     requestor.stop()