TrixyZ/trainer/components/callbacks/manage_data_callbacks.py

from dash.dependencies import Input, Output, State, ALL
from dash import html, callback_context
from utils.file_utils import list_files_in_directory, save_file, delete_file, is_valid_wav
from components.hotwords import generate_training_data_table
import os
import base64

TRAINING_HOTWORD_DIR = './training/hotword/'

def register_manage_data_callbacks(app):

    @app.callback(
        Output('file-preview-list', 'children'),
        [
            Input('upload-training-data', 'contents'),
            Input({'type': 'delete-preview-file', 'index': ALL}, 'n_clicks')
        ],
        [
            State('upload-training-data', 'filename'),
            State('upload-training-data', 'last_modified'),
            State('file-preview-list', 'children')
        ]
    )
    def manage_file_previews(contents, n_clicks, filenames, last_modified, preview_list):
        ctx = callback_context

        if not ctx.triggered:
            return preview_list

        if preview_list is None:
            preview_list = []

        triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]

        if 'upload-training-data' in triggered_id and contents:
            files = []
            for content, filename in zip(contents, filenames):
                files.append(html.Div([
                    html.Audio(src=f'data:audio/wav;base64,{content.split(",")[1]}', controls=True),
                    html.Button("Delete", id={'type': 'delete-preview-file', 'index': filename}),
                    html.Div(content, style={'display': 'none'}, id={'type': 'file-content', 'index': filename})
                ], id={'type': 'preview-file', 'index': filename}))
            return preview_list + files

        if 'delete-preview-file' in triggered_id:
            filename_to_delete = ctx.triggered[0]['prop_id'].split('.')[1]
            new_preview_list = [
                item for item in preview_list if item['props']['id']['index'] != filename_to_delete
            ]
            return new_preview_list

        return preview_list

    @app.callback(
        Output('upload-status', 'children'),
        Input('confirm-upload', 'n_clicks'),
        State('file-preview-list', 'children'),
        State('url', 'pathname')
    )
    def upload_files(n_clicks, preview_list, pathname):
        if n_clicks and preview_list:
            model_name = pathname.split('/')[-1]
            for item in preview_list:
                content = item['props']['children'][2]['props']['children']
                filename = item['props']['id']['index']
                data = content.split(",")[1]
                file_path = f'{TRAINING_HOTWORD_DIR}/{model_name}/{filename}'
                save_file(file_path, base64.b64decode(data))
                if not is_valid_wav(file_path):
                    os.remove(file_path)
                    return f"File {filename} is not a valid WAV file.", True
            training_data = list_files_in_directory(f'{TRAINING_HOTWORD_DIR}/{model_name}')
            return generate_training_data_table(model_name, training_data)
        return html.Div()

    @app.callback(
        Output('hotword-manage-training-data', 'children'),
        [
            Input({"type": "delete-button", "index": ALL}, 'n_clicks'),
            Input({"type": "toggle-button", 'index': ALL}, 'n_clicks')
        ],
        [
            State({"type": "delete-button", "index": ALL}, 'id'),
            State({"type": "toggle-button", "index": ALL}, 'id'),
            State('url', 'pathname')
        ]
    )
    def manage_training_data(n_clicks_delete, n_clicks_toggle, delete_ids, toggle_ids, pathname):
        ctx = callback_context

        if not ctx.triggered:
            return html.Div()

        model_name = pathname.split('/')[-1]
        triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]

        if 'delete-button' in triggered_id:
            for item in delete_ids:
                if item['id'] == triggered_id:
                    file_path = os.path.join(TRAINING_HOTWORD_DIR, model_name, item['index'])
                    delete_file(file_path)
            training_data = list_files_in_directory(f'{TRAINING_HOTWORD_DIR}/{model_name}')
            return generate_training_data_table(model_name, training_data)

        if 'toggle-button' in triggered_id:
            for item in toggle_ids:
                if item['id'] == triggered_id:
                    file_path = os.path.join(TRAINING_HOTWORD_DIR, model_name, item['index'])
                    new_file_path = os.path.join(TRAINING_HOTWORD_DIR, model_name, '_' + item['index'] if not item['index'].startswith('_') else item['index'][1:])
                    os.rename(file_path, new_file_path)
            training_data = list_files_in_directory(f'{TRAINING_HOTWORD_DIR}/{model_name}')
            return generate_training_data_table(model_name, training_data)

        return html.Div()