TrixyZ/trainer/app.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import dash
import dash_bootstrap_components as dbc
from dash import html, dcc
from dash.dependencies import Input, Output
from components.navigation import create_navbar
from components.hotwords import HotwordsLayout, ManageHotwordLayout
from components.intents import IntentsLayout
from components.nlp import NLPLayout
from components.voice_detection import VoiceDetectionLayout
from components.tts import TTSLayout
from components.data_source import DataSourceLayout
from components.callbacks import register_callbacks
from flask import send_from_directory, send_file
from utils.file_utils import ensure_directory_exists
from importlib import import_module
from utils.training_interface import TrainerInterface
from utils.file_utils import trainer_classes

# Ensure directories exist
ensure_directory_exists('./static/css/')
ensure_directory_exists('./trainer_scripts/')


def load_trainer_classes():
    trainer_dir = './trainer_scripts'
    for file in os.listdir(trainer_dir):
        if file.endswith('.py'):
            module_name = file[:-3]
            module = import_module(f'trainer_scripts.{module_name}')
            for attr in dir(module):
                cls = getattr(module, attr)
                if isinstance(cls, type) and issubclass(cls, TrainerInterface) and cls is not TrainerInterface:
                    trainer_classes[module_name] = cls

load_trainer_classes()

app = dash.Dash(__name__, external_stylesheets=['/static/css/bootstrap.min.css'], suppress_callback_exceptions=True)
app.title = "Trainer"

app.layout = dbc.Container(
    [
        dcc.Location(id='url', refresh=False),
        create_navbar(),
        html.Div(id='main-navigation', className='mt-4'),
    ],
    fluid=True,
)

@app.server.route('/static/<path:path>')
def serve_static(path):
    return send_from_directory('static', path)

@app.server.route('/audio/<path:filename>')
def serve_audio(filename: str):
    ext = filename.split(".")[-1]
    if ext == "png":
        return display_image(filename[0:-4], "png")
    elif ext == "jpg":
        return display_image(filename[0:-4], "jpg")
    elif ext == "jpeg":
        return display_image(filename[0:-5], "jpg")
    else:
        return send_from_directory('training', filename)

def display_image(filename, ext):
    import matplotlib.pyplot as plt
    import numpy as np
    import wave

    wavefile = f"training/{filename}.wav"
    imagefile = f"training/{filename}.{ext}"

    if not os.path.exists(imagefile):
        print(f"Generate Wave Image: {filename} as {ext} from {wavefile}")
        with wave.open(wavefile, 'r') as wav_file:
            signal = wav_file.readframes(-1)
            signal = np.frombuffer(signal, dtype=np.int16)
            channels = [[] for channel in range(wav_file.getnchannels())]
            for index, datum in enumerate(signal):
                channels[index % len(channels)].append(datum)
            fs = wav_file.getframerate()
            Time = np.linspace(0, len(signal) / len(channels) / fs, num=int(len(signal) / len(channels)))
            
            plt.figure(figsize=(8, 1))
            for channel in channels:
                plt.plot(Time, channel)
            plt.xlabel("")
            plt.ylabel("")
            plt.title("")
            plt.grid()
            plt.tight_layout()
            plt.axis('off')
            if ext == "png":
                plt.savefig(imagefile, format='png', transparent=True, bbox_inches='tight', pad_inches=0)
            else:
                plt.savefig(imagefile, format='jpg', bbox_inches='tight', pad_inches=0)
            plt.close()
    
    if ext == "png":
        return send_file(imagefile, mimetype='image/png')
    else:
        return send_file(imagefile, mimetype='image/jpeg')

@app.callback(Output('main-navigation', 'children'),
              [Input('url', 'pathname')])
def display_page(pathname):
    if pathname == '/hotwords':
        return HotwordsLayout()
    elif pathname == '/intents':
        return IntentsLayout()
    elif pathname == '/nlp':
        return NLPLayout()
    elif pathname == '/voice_detection':
        return VoiceDetectionLayout()
    elif pathname == '/tts':
        return TTSLayout()
    elif pathname == '/data_source':
        return DataSourceLayout()
    elif pathname and pathname.startswith('/manage_hotword/'):
        model_name = pathname.split('/')[-1]
        return ManageHotwordLayout(model_name)
    else:
        return html.H1("Home", className="text-center")

register_callbacks(app)

if __name__ == "__main__":
    app.run_server(debug=True)