import os import logging import time import numpy as np import pandas as pd from sklearn.model_selection import train_test_split from tensorflow.keras.models import Sequential, load_model from tensorflow.keras.layers import Dense, Activation, Dropout from tensorflow.keras.utils import to_categorical from utils.training_interface import TrainerInterface from utils.file_utils import list_files_in_directory import json # Define global variables for tracking progress current_progress = 0 current_loss = 0.0 current_accuracy = 0.0 total_epochs = 0 total_batches = 0 class TensorFlowTrainer(TrainerInterface): def __init__(self, learning_rate=0.001): self.model_path = None self.training_directories = [] self.sample_directory = None self.background_directory = None self.dropout_rate = 0.5 self.epochs = 10 self.batch_size = 32 self.learning_rate = learning_rate self.model = None self.current_epoch = 0 self.stop_training = False self.training_status = "Not started" self.train_data = None self.val_data = None def set_model_path(self, model_path): self.model_path = model_path self.model = self.load_model() def set_training_directories(self, directories): self.training_directories = directories def set_sample_directory(self, directory): self.sample_directory = directory def set_background_directory(self, directory): self.background_directory = directory def set_dropout_rate(self, rate): self.dropout_rate = rate def set_epochs(self, epochs): self.epochs = epochs def set_batch_size(self, batch_size): self.batch_size = batch_size def get_training_data(self): return self.training_directories def get_model_path(self): return self.model_path def get_epochs(self): return self.epochs def get_batch_size(self): return self.batch_size def get_dropout_rate(self): return self.dropout_rate def get_training_directories(self): return self.training_directories def get_sample_directory(self): return self.sample_directory def get_background_directory(self): return self.background_directory def set_training_data(self, features, labels): X_train, X_val, y_train, y_val = train_test_split(features, labels, test_size=0.2, random_state=42) self.train_data = (X_train, y_train) self.val_data = (X_val, y_val) def load_model(self): if os.path.exists(self.model_path + "hotword_model.h5"): return load_model(self.model_path + "hotword_model.h5") else: model = Sequential([ Dense(256, input_shape=(40,)), Activation('relu'), Dropout(self.dropout_rate), Dense(256), Activation('relu'), Dropout(self.dropout_rate), Dense(len(self.training_directories) + 1, activation='softmax') ]) model.compile( loss="categorical_crossentropy", optimizer='adam', metrics=['accuracy'] ) return model def start(self): self.training_status = "Training" self.stop_training = False self.run_training() def start_new(self): self.training_status = "Training" self.stop_training = False self.model = self.load_model() self.current_epoch = 0 self.run_training() def pause(self): self.training_status = "Paused" self.stop_training = True def resume(self): self.training_status = "Training" self.stop_training = False self.run_training() def stop(self): self.training_status = "Stopped" self.stop_training = True def preprocess_training_data(self): import librosa data_path_dict = {} data_path_dict[0] = list_files_in_directory(self.background_directory, extension='.wav') for idx, directory in enumerate(self.training_directories): data_path_dict[idx+1] = list_files_in_directory(directory, extension='.wav') all_data = [] total_files = sum(len(files) for files in data_path_dict.values()) processed_files = 0 for class_label, list_of_files in data_path_dict.items(): for single_file in list_of_files: try: audio, sample_rate = librosa.load(single_file) mfcc = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40) mfcc_processed = np.mean(mfcc.T, axis=0) all_data.append([mfcc_processed, class_label]) processed_files += 1 global current_progress current_progress = (processed_files / total_files) * 100 except Exception as e: print(f"Exception: {str(e)}") df = pd.DataFrame(all_data, columns=["feature", "class_label"]) df.to_pickle(self.model_path + "audio_data.csv") def run_training(self): global current_progress, current_loss, current_accuracy, total_epochs, total_batches self.preprocess_training_data() start_time = time.time() df = pd.read_pickle(self.model_path + "audio_data.csv") X = df["feature"].values X = np.concatenate(X, axis=0).reshape(len(X), 40) y = np.array(df["class_label"].tolist()) y = to_categorical(y) self.set_training_data(X, y) X_train, y_train = self.train_data total_batches = int(np.ceil(len(X_train) / self.batch_size)) total_epochs = self.epochs training_info = { "trainer_script": "tensorflow_trainer", "training_time": 0, "model_size": 0, "epochs": self.epochs, "dropout_rate": self.dropout_rate, "batch_size": self.batch_size, "total_training_data": len(X), "training_directories": self.training_directories, "accuracy": [], "loss": [], "test_accuracy": 0 } for epoch in range(self.current_epoch, self.epochs): if self.stop_training: break self.current_epoch = epoch history = self.model.fit(X_train, y_train, epochs=1, batch_size=self.batch_size) current_loss = history.history['loss'][-1] current_accuracy = history.history['accuracy'][-1] current_progress = (epoch + 1) / self.epochs * 100 training_info["accuracy"].append(current_accuracy) training_info["loss"].append(current_loss) print(f"Epoch {epoch + 1}/{self.epochs}, Loss: {current_loss}, Accuracy: {current_accuracy}") epoch_time = time.time() - start_time training_info["training_time"] = epoch_time if not self.stop_training: self.model.save(self.model_path + "hotword_model.h5") self.evaluate_model(X, y) training_info["test_accuracy"] = current_accuracy training_info["model_size"] = os.path.getsize(self.model_path + "hotword_model.h5") / 1024 # in KB self.save_training_info(training_info) self.training_status = "Finished" else: self.training_status = "Paused" def evaluate_model(self, X_test, y_test): loss, accuracy = self.model.evaluate(X_test, y_test) print(f"Test Accuracy: {accuracy * 100:.2f}%") def save_training_info(self, info): with open(self.model_path + "info_tensorflow_trainer.json", "w") as f: json.dump(info, f) def get_training_info(self): return { "status": self.training_status, "current_epoch": self.current_epoch, "total_epochs": self.epochs, "loss": current_loss, "accuracy": current_accuracy }