Classification improvements (#19020)

* Move classification training to full process

* Sort class images

frigate/util/classification.py

@@ -20,106 +20,117 @@ LEARNING_RATE = 0.001
 logger = logging.getLogger(__name__)
 
 
-def __generate_representative_dataset_factory(dataset_dir: str):
+class ClassificationTrainingProcess(FrigateProcess):
-    def generate_representative_dataset():
+    def __init__(self, model_name: str) -> None:
-        image_paths = []
+        super().__init__(
-        for root, dirs, files in os.walk(dataset_dir):
+            stop_event=None,
-            for file in files:
+            name=f"model_training:{model_name}",
-                if file.lower().endswith((".jpg", ".jpeg", ".png")):
+        )
-                    image_paths.append(os.path.join(root, file))
+        self.model_name = model_name
 
-        for path in image_paths[:300]:
+    def run(self) -> None:
-            img = cv2.imread(path)
+        self.pre_run_setup()
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        self.__train_classification_model()
-            img = cv2.resize(img, (224, 224))
-            img_array = np.array(img, dtype=np.float32) / 255.0
-            img_array = img_array[None, ...]
-            yield [img_array]
 
-    return generate_representative_dataset
+    def __generate_representative_dataset_factory(self, dataset_dir: str):
+        def generate_representative_dataset():
+            image_paths = []
+            for root, dirs, files in os.walk(dataset_dir):
+                for file in files:
+                    if file.lower().endswith((".jpg", ".jpeg", ".png")):
+                        image_paths.append(os.path.join(root, file))
 
+            for path in image_paths[:300]:
+                img = cv2.imread(path)
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                img = cv2.resize(img, (224, 224))
+                img_array = np.array(img, dtype=np.float32) / 255.0
+                img_array = img_array[None, ...]
+                yield [img_array]
 
-@redirect_output_to_logger(logger, logging.DEBUG)
+        return generate_representative_dataset
-def __train_classification_model(model_name: str) -> bool:
-    """Train a classification model."""
 
-    # import in the function so that tensorflow is not initialized multiple times
+    @redirect_output_to_logger(logger, logging.DEBUG)
-    import tensorflow as tf
+    def __train_classification_model(self) -> bool:
-    from tensorflow.keras import layers, models, optimizers
+        """Train a classification model."""
-    from tensorflow.keras.applications import MobileNetV2
-    from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
-    logger.info(f"Kicking off classification training for {model_name}.")
+        # import in the function so that tensorflow is not initialized multiple times
-    dataset_dir = os.path.join(CLIPS_DIR, model_name, "dataset")
+        import tensorflow as tf
-    model_dir = os.path.join(MODEL_CACHE_DIR, model_name)
+        from tensorflow.keras import layers, models, optimizers
-    num_classes = len(
+        from tensorflow.keras.applications import MobileNetV2
-        [
+        from tensorflow.keras.preprocessing.image import ImageDataGenerator
-            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
+        logger.info(f"Kicking off classification training for {self.model_name}.")
-    base_model = MobileNetV2(
+        dataset_dir = os.path.join(CLIPS_DIR, self.model_name, "dataset")
-        input_shape=(224, 224, 3),
+        model_dir = os.path.join(MODEL_CACHE_DIR, self.model_name)
-        include_top=False,
+        num_classes = len(
-        weights="imagenet",
+            [
-        alpha=0.35,
+                d
-    )
+                for d in os.listdir(dataset_dir)
-    base_model.trainable = False  # Freeze pre-trained layers
+                if os.path.isdir(os.path.join(dataset_dir, d))
+            ]
+        )
 
-    model = models.Sequential(
+        # Start with imagenet base model with 35% of channels in each layer
-        [
+        base_model = MobileNetV2(
-            base_model,
+            input_shape=(224, 224, 3),
-            layers.GlobalAveragePooling2D(),
+            include_top=False,
-            layers.Dense(128, activation="relu"),
+            weights="imagenet",
-            layers.Dropout(0.3),
+            alpha=0.35,
-            layers.Dense(num_classes, activation="softmax"),
+        )
-        ]
+        base_model.trainable = False  # Freeze pre-trained layers
-    )
 
-    model.compile(
+        model = models.Sequential(
-        optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
+            [
-        loss="categorical_crossentropy",
+                base_model,
-        metrics=["accuracy"],
+                layers.GlobalAveragePooling2D(),
-    )
+                layers.Dense(128, activation="relu"),
+                layers.Dropout(0.3),
+                layers.Dense(num_classes, activation="softmax"),
+            ]
+        )
 
-    # create training set
+        model.compile(
-    datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
+            optimizer=optimizers.Adam(learning_rate=LEARNING_RATE),
-    train_gen = datagen.flow_from_directory(
+            loss="categorical_crossentropy",
-        dataset_dir,
+            metrics=["accuracy"],
-        target_size=(224, 224),
+        )
-        batch_size=BATCH_SIZE,
-        class_mode="categorical",
-        subset="training",
-    )
 
-    # write labelmap
+        # create training set
-    class_indices = train_gen.class_indices
+        datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
-    index_to_class = {v: k for k, v in class_indices.items()}
+        train_gen = datagen.flow_from_directory(
-    sorted_classes = [index_to_class[i] for i in range(len(index_to_class))]
+            dataset_dir,
-    with open(os.path.join(model_dir, "labelmap.txt"), "w") as f:
+            target_size=(224, 224),
-        for class_name in sorted_classes:
+            batch_size=BATCH_SIZE,
-            f.write(f"{class_name}\n")
+            class_mode="categorical",
+            subset="training",
+        )
 
-    # train the model
+        # write labelmap
-    model.fit(train_gen, epochs=EPOCHS, verbose=0)
+        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")
 
-    # convert model to tflite
+        # train the model
-    converter = tf.lite.TFLiteConverter.from_keras_model(model)
+        model.fit(train_gen, epochs=EPOCHS, verbose=0)
-    converter.optimizations = [tf.lite.Optimize.DEFAULT]
-    converter.representative_dataset = __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
+        # convert model to tflite
-    with open(os.path.join(model_dir, "model.tflite"), "wb") as f:
+        converter = tf.lite.TFLiteConverter.from_keras_model(model)
-        f.write(tflite_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
+        with open(os.path.join(model_dir, "model.tflite"), "wb") as f:
+            f.write(tflite_model)
 
 
 @staticmethod
@@ -138,12 +149,7 @@ def kickoff_model_training(
     # run training in sub process so that
     # tensorflow will free CPU / GPU memory
     # upon training completion
-    training_process = FrigateProcess(
+    training_process = ClassificationTrainingProcess(model_name)
-        None,
-        target=__train_classification_model,
-        name=f"model_training:{model_name}",
-        args=(model_name,),
-    )
     training_process.start()
     training_process.join()
 

web/src/views/classification/ModelTrainingView.tsx

@@ -577,9 +577,14 @@ function DatasetGrid({
 }: DatasetGridProps) {
   const { t } = useTranslation(["views/classificationModel"]);
 
+  const classData = useMemo(
+    () => images.sort((a, b) => a.localeCompare(b)),
+    [images],
+  );
+
   return (
     <div className="flex flex-wrap gap-2 overflow-y-auto p-2">
-      {images.map((image) => (
+      {classData.map((image) => (
         <div
           className={cn(
             "flex w-60 cursor-pointer flex-col gap-2 rounded-lg bg-card outline outline-[3px]",