diff --git a/frigate/util/classification.py b/frigate/util/classification.py index 6eab829f2..533c1345a 100644 --- a/frigate/util/classification.py +++ b/frigate/util/classification.py @@ -20,106 +20,117 @@ LEARNING_RATE = 0.001 logger = logging.getLogger(__name__) -def __generate_representative_dataset_factory(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)) +class ClassificationTrainingProcess(FrigateProcess): + def __init__(self, model_name: str) -> None: + super().__init__( + stop_event=None, + name=f"model_training:{model_name}", + ) + self.model_name = model_name - 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] + def run(self) -> None: + self.pre_run_setup() + self.__train_classification_model() - 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) -def __train_classification_model(model_name: str) -> bool: - """Train a classification model.""" + return generate_representative_dataset - # 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 + @redirect_output_to_logger(logger, logging.DEBUG) + def __train_classification_model(self) -> bool: + """Train a classification model.""" - logger.info(f"Kicking off classification training for {model_name}.") - dataset_dir = os.path.join(CLIPS_DIR, model_name, "dataset") - model_dir = os.path.join(MODEL_CACHE_DIR, model_name) - num_classes = len( - [ - d - for d in os.listdir(dataset_dir) - if os.path.isdir(os.path.join(dataset_dir, d)) - ] - ) + # 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 - # 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 + 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) + num_classes = len( + [ + d + for d in os.listdir(dataset_dir) + if os.path.isdir(os.path.join(dataset_dir, d)) + ] + ) - 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) + # 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") - # convert model to tflite - converter = tf.lite.TFLiteConverter.from_keras_model(model) - 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() + # train the model + model.fit(train_gen, epochs=EPOCHS, verbose=0) - # write model - with open(os.path.join(model_dir, "model.tflite"), "wb") as f: - f.write(tflite_model) + # 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 + 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( - None, - target=__train_classification_model, - name=f"model_training:{model_name}", - args=(model_name,), - ) + training_process = ClassificationTrainingProcess(model_name) training_process.start() training_process.join() diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx index 14de1a118..145004ec3 100644 --- a/web/src/views/classification/ModelTrainingView.tsx +++ b/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 (
- {images.map((image) => ( + {classData.map((image) => (