TrixyZ/trainer/trainer_scripts/tensorflow_trainer.py

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
        }