TrixyZ/trainer/components/hotwords.py

from dash import html, dcc, callback_context, dash_table
import dash_bootstrap_components as dbc
from dash.dependencies import Input, Output, State
from utils.file_utils import list_directories_in_directory, list_files_in_directory, get_file_metadata, get_training_time, get_file_size, trainer_classes, get_Size_String
import os
import json
import plotly.graph_objs as go

TRAINING_HOTWORD_DIR = './training/hotword/'

def HotwordsLayout():
    hotword_models = list_directories_in_directory(TRAINING_HOTWORD_DIR)
    model_table = generate_model_table(hotword_models)

    return html.Div([
        html.H2("Hotword Models"),
        dbc.Button("Add New Hotword", id="add-new-hotword", color="primary", className="mb-3"),
        html.Div(id="new-hotword-input"),
        model_table,
        dcc.Store(id='hidden-model-name'),
        html.Div(id="hotword-manage-training-data")
    ])

def generate_model_table(models):
    header = [
        html.Thead(html.Tr([html.Th("Model Name"), html.Th("Actions")]))
    ]
    rows = []
    for model in models:
        rows.append(html.Tr([
            html.Td(model),
            html.Td([
                dbc.Button("Details", href=f"/manage_hotword/{model}", color='primary', className='mr-2'),
                dbc.Button("Delete", id={'type': 'delete-hotword', 'index': model}, color='danger')
            ])
        ]))
    return dbc.Table(header + [html.Tbody(rows)], bordered=True, hover=True)

def ManageHotwordLayout(model_name):
    page_size = 10
    training_data = list_files_in_directory(f"{TRAINING_HOTWORD_DIR}/{model_name}/", ".wav")
    training_data_table, row_count = generate_training_data_table(model_name, training_data, 1, page_size)
    trainer_options = [{'label': name, 'value': name} for name in trainer_classes.keys()]

    training_time = get_training_time(model_name)
    
    tab_trainer = []
    for trainer_info in trainer_options:
        infoFile = f"./model/hotword/{model_name}/info_{trainer_info['value']}.json"
        if os.path.exists(infoFile):
            with open(infoFile, 'r') as file:
                info = json.load(file)
                epochs = '-'
                batch_size = '-'
                dropout_rate = '-'
                training_method = '-'
                duration = '-'
                testAccuracy = '-'
                accu = []
                loss = []
                size = '-'
                if 'epochs' in info:
                    epochs = info['epochs']
                if 'batch_size' in info:
                    batch_size = info['batch_size']
                if 'dropout_rate' in info:
                    dropout_rate = info['dropout_rate']
                if 'trainer_script' in info:
                    training_method = info['trainer_script']
                if 'accuracy' in info:
                    accu = info['accuracy']
                if 'loss' in info:
                    loss = info['loss']
                if 'training_time' in info:
                    duration = info['training_time']
                if 'test_accuracy' in info:
                    testAccuracy = info['test_accuracy']
                if 'model_size' in info:
                    size = info['model_size']
                chart_fig1 = go.Figure(data=[go.Scatter(y=accu)])
                chart_fig2 = go.Figure(data=[go.Scatter(y=loss)])
                entries = [
                    dcc.Graph(figure=chart_fig1),
                    dcc.Graph(figure=chart_fig2),
                    html.P("Epochen: "+str(epochs), className="card-text"),
                    html.P("Batch Size: "+str(batch_size), className="card-text"),
                    html.P("Dropout Rate: "+str(dropout_rate), className="card-text"),
                    html.P("Methode: "+str(training_method), className="card-text"),
                    html.P("Duration: "+str(duration), className="card-text"),
                    html.P("Test Accuracy: "+str(testAccuracy), className="card-text"),
                    html.P("Size: "+get_Size_String(size*1024), className="card-text"),
                ]
                tab_trainer.append(
                    dbc.Tab(entries
                    , label=trainer_info['label'], tab_id=trainer_info['value'])
                )
        else:
            tab_trainer.append(
                dbc.Tab("",label=trainer_info['label'], tab_id=trainer_info['value'], disabled=True)
            )
    
    # Pagination logic
    num_pages = (row_count // page_size) + (1 if row_count % page_size > 0 else 0)
    pagination = dbc.Pagination(id='data-table-pagination', max_value=num_pages, active_page=1)

    return html.Div([
        html.H2(f"Manage Hotword Model: {model_name}"),
        dbc.Card([dbc.CardHeader("Model Infos"),dbc.CardBody(dbc.Tabs(tab_trainer))]),
        dbc.Card([
            dbc.CardHeader("Training Controls"),
            dbc.CardBody([
                #dbc.Row([
                    #dbc.Col(dbc.Button("Start Training", id="start-training-hotword", color="success", className="mb-3")),
                    #dbc.Col(dbc.Button("Optimize Training", id="optimize-training-hotword", color="primary", className="mb-3 ml-3")),
                    #dbc.Col(dbc.Button("Stop Training", id="stop-training-hotword", color="danger", className="mb-3 ml-3")),
                    dbc.ButtonGroup([
                        dbc.DropdownMenu(
                            label="Start training",
                            children=[
                                dbc.DropdownMenuItem("Item 1"),
                                dbc.DropdownMenuItem("Item 2"),
                                dbc.DropdownMenuItem("Item 3"),
                            ],
                        ),
                        dbc.Button("Start Training", id="start-training-hotword", color="success", className="mb-3"),
                        dbc.Button("Optimize Training", id="optimize-training-hotword", color="primary", className="mb-3 ml-3"),
                        dbc.Button("Stop Training", id="stop-training-hotword", color="danger", className="mb-3 ml-3"),
                    ]),
                #]),
                dbc.Row([
                    dbc.Col(
                        dcc.Dropdown(
                            id="training-method-dropdown",
                            options=trainer_options,
                            value=trainer_options[0]["value"],  # default value
                            clearable=False,
                            className="mb-3",
                            placeholder="Select Training Method",
                        )
                    ),
                    dbc.Col(
                        dcc.Input(
                            id="dropout-rate",
                            type="number",
                            value=0.5,
                            step=0.1,
                            min=0,
                            max=1,
                            className="mb-3",
                            placeholder="Dropout Rate"
                        )
                    ),
                    dbc.Col(
                        dcc.Slider(0, 10000, value=3000, id="num-epochs",tooltip={"placement": "bottom", "always_visible": False})
                        #dcc.Input(
                        #    id="num-epochs",
                        #    type="number",
                        #    value=10,
                        #    step=1,
                        #    min=1,
                        #    className="mb-3",
                        #    placeholder="Epochs"
                        #)
                    ),
                    dbc.Col(
                        dcc.Slider(0, 256, marks={
                            8: '8',
                            16: '16',
                            32: '32',
                            64: '64',
                            128: '128',
                            256: '256',
                        }, value=32, id="batch-size", step=None)
                        #dcc.Input(
                        #    id="batch-size",
                        #    type="number",
                        #    value=32,
                        #    step=1,
                        #    min=1,
                        #    className="mb-3",
                        #    placeholder="Batch Size"
                        #)
                    )
                ]),
                html.Div(id="training-status", className='mt-3'),
                html.Div(id="training-timer", className='mt-3'),
                dbc.Card([
                    dbc.CardHeader("Training Progress"),
                    dbc.CardBody([
                        dcc.Graph(id="training-chart", figure=[]),
                        dbc.Progress(id="training-progress", striped=True, animated=True, style={"height": "20px"}),
                    ])
                ]),
            ])
        ]),
        dbc.Card([
            dbc.CardHeader("Upload Training Data"),
            dbc.CardBody([
                dcc.Upload(
                    id="upload-training-data",
                    children=html.Div([
                        'Drag and Drop or ',
                        html.A('Select Files')
                    ]),
                    style={
                        'width': '100%',
                        'height': '60px',
                        'lineHeight': '60px',
                        'borderWidth': '1px',
                        'borderStyle': 'dashed',
                        'borderRadius': '5px',
                        'textAlign': 'center',
                        'margin': '10px'
                    },
                    multiple=True
                ),
                html.Div(id="file-preview-list"),
                dbc.Button("Upload Files", id="confirm-upload", color='success', className='mt-3'),
                html.Div(id="upload-status"),
            ])
        ]),
        dbc.Card([
            dbc.CardHeader("Training Data"),
            dbc.CardBody([
                training_data_table, pagination
            ])
        ])
    ])

def generate_training_data_table(model_name, files, page=1, page_size=10):
    import os
    header = [
        html.Thead(html.Tr([html.Th("Audio"), html.Th("File Name"), html.Th("Display"), html.Th("Length (s)"), html.Th("Size (KB)"), html.Th("Actions")]))
    ]
    rows = []
    for file in files:
        filex = os.path.basename(file)
        if not os.path.isdir(f"{TRAINING_HOTWORD_DIR}/{model_name}/{filex}"):
            file_name = filex.lstrip('_')
            file_path = f"{TRAINING_HOTWORD_DIR}/{model_name}/{filex}"
            fileImg = filex[0:-3] + "png"
            try:
                length, size = get_file_metadata(file_path)
                rows.append(html.Tr([
                    html.Td(html.Audio(src=f"/audio/hotword/{model_name}/{filex}", controls=True)),
                    html.Td(file_name),
                    html.Td(html.Img(src=f"/audio/hotword/{model_name}/{fileImg}")),
                    html.Td(length),
                    html.Td(size),
                    html.Td([
                        dbc.Button("Delete", id={'type': 'delete-button', 'index': file}, color='danger', className='mr-2'),
                        dbc.Button("Deactivate" if not filex.startswith('_') else "Activate", id={'type': 'toggle-button', 'index': filex}, color='warning')
                    ])
                ]))
            except ValueError as e:
                rows.append(html.Tr([
                    html.Td("-"),
                    html.Td(file_name),
                    html.Td("-"),
                    html.Td("-"),
                    html.Td("-"),
                    html.Td([
                        dbc.Button("Delete", id={'type': 'delete-button', 'index': file}, color='danger', className='mr-2')
                    ])
                ]))
            except FileNotFoundError as e:
                rows.append(html.Tr([
                    html.Td("-"),
                    html.Td(file_name),
                    html.Td("-"),
                    html.Td("-"),
                    html.Td("-"),
                    html.Td("File not found")
                ]))

    start = (page - 1) * 10
    end = start + 10
    row_count = len(rows)
    table_body = html.Tbody(id='data-table-body', children=rows[start:end])

    return html.Div([
        dbc.Table(header + [table_body], bordered=True, hover=True)
    ]), row_count