TrixyZ/trainer/utils/file_utils.py

from cmath import e
import os
from re import A
import wave
import logging
from pydub import AudioSegment
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import librosa.display

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
trainer_classes = {}

def ensure_directory_exists(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

def list_files_in_directory(directory, extension=None):
    if not os.path.exists(directory):
        raise FileNotFoundError(f"Directory {directory} does not exist.")
    
    files = []
    for file in os.listdir(directory):
        if extension and not file.endswith(extension):
            continue
        files.append(directory+file)
    return files

def list_directories_in_directory(directory):
    if not os.path.exists(directory):
        raise FileNotFoundError(f"Directory {directory} does not exist.")
    
    directories = [d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))]
    return directories

def save_file(directory, filename, data):
    ensure_directory_exists(directory)
    file_path = os.path.join(directory, filename)
    with open(file_path, 'wb') as file:
        file.write(data)

def delete_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)

def get_file_metadata(file_path):
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File {file_path} does not exist.")
    
    try:
        with wave.open(file_path, 'rb') as wf:
            frames = wf.getnframes()
            rate = wf.getframerate()
            length = frames / float(rate)
        size = os.path.getsize(file_path) / 1024  # size in KB
        return round(length, 2), round(size, 2)
    except wave.Error as e:
        logger.error(f"Error processing WAV file {file_path}: {e}")
        raise ValueError(f"Error processing WAV file: {e}")

def is_valid_wav(file_path):
    try:
        with wave.open(file_path, 'rb') as wf:
            return True
    except wave.Error:
        return False

def get_file_size(file_path):
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File {file_path} does not exist.")
    size = os.path.getsize(file_path)
    return get_Size_String(size)

def get_Size_String(value):
    ext = ["Byte","KB","MB","GB","TB"]
    i = 0
    while value > 1024:
        value /= 1024
        i += 1
    return f"{value:.2f} {ext[i]}"

def update_training_time(model_name, additional_time):
    training_time_path = f"./model/hotword/{model_name}/training_time.txt"
    ensure_directory_exists(os.path.dirname(training_time_path))
    if os.path.exists(training_time_path):
        with open(training_time_path, "r") as f:
            total_time = int(f.read().strip())
    else:
        total_time = 0
    total_time += additional_time
    with open(training_time_path, "w") as f:
        f.write(str(total_time))

def get_training_time(model_name):
    training_time_path = f"./model/hotword/{model_name}/training_time.txt"
    if os.path.exists(training_time_path):
        with open(training_time_path, "r") as f:
            return int(f.read().strip())
    return 0

def resample_training_data(training_directory, background_directory=None):
    print("  training_directory: "+training_directory)
    print("  background_directory: "+background_directory)
    input_dir = os.path.join(training_directory)
    output_dir = os.path.join(input_dir, 'processed/')
    print("  input_dir: "+input_dir)
    print("  output_dir: "+output_dir)
    ensure_directory_exists(output_dir)

    hotword_files = list_files_in_directory(input_dir, extension='.wav')
    if background_directory:
        background_files = list_files_in_directory(background_directory, extension='.wav')
    else:
        background_files = []

    for hotword_file in hotword_files:
        hotword_audio = AudioSegment.from_wav(hotword_file)
        fname1=os.path.splitext(os.path.basename(hotword_file))[0]+".wav"
        hotword_audio.export(os.path.join(output_dir, fname1), format='wav')  # Save raw hotword

        for background_file in background_files:
            fname1=os.path.splitext(os.path.basename(hotword_file))[0]
            fname2=os.path.splitext(os.path.basename(background_file))[0]
            mixed_filename = f"{fname1}_{fname2}.wav"
            ofile = os.path.join(output_dir, mixed_filename)
            if os.path.exists(ofile)==False:
                background_audio = AudioSegment.from_wav(background_file)
                mixed_audio = hotword_audio.overlay(background_audio, position=0)
                print(f"  Mixed Filename: {mixed_filename}")
                mixed_audio.export(ofile, format='wav')

def preprocessing_training_data(model_dir,data_dir, sc_data_dir, fail_data_dir):
    import librosa
    
    data_path_dict = {
        0: list_files_in_directory(fail_data_dir, extension='.wav'),
        1: list_files_in_directory(data_dir, extension='.wav'),
        2: list_files_in_directory(sc_data_dir, extension='.wav')
    }

    #walley_sample = "training/background/rec04.wav"
    #data, sample_rate = librosa.load(walley_sample)
    #plt.title("Wave Form")
    #librosa.display.waveplot(data, sr=sample_rate)
    #plt.show()
    #mfccs = librosa.feature.mfcc(y=data, sr=sample_rate, n_mfcc=40)
    #print("Shape of mfcc:", mfccs.shape)
    #plt.title("MFCC")
    #librosa.display.specshow(mfccs, sr=sample_rate, x_axis='time')
    #plt.show()

    all_data = []
    print("  Load training data")
    for class_label, list_of_files in data_path_dict.items():
        for single_file in list_of_files:
            try:
                if single_file[0:2]=="./":
                    single_file = single_file[2:]
                print(f"    {class_label}: {single_file}")
                audio, sample_rate = librosa.load(single_file)
                print(f"      Sample-Rate: {sample_rate}")
                mfcc = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
                print(f"      mfcc: {mfcc}")
                mfcc_processed = np.mean(mfcc.T, axis=0)
                print(f"      mfcc_processed: mfcc_processed")
                all_data.append([mfcc_processed, class_label])
                print("      Next")
            except Exception as e:
                print(f"      Exception: {str(e)}")
        print(f"  Info: Succesfully Preprocessed Class Label {class_label}")

    df = pd.DataFrame(all_data, columns=["feature", "class_label"])
    df.to_pickle(model_dir+"/audio_data.csv")