app.py

import random
import Phonemize
from Levenshtein import editops
from gradio.components import Audio, Dropdown, Textbox, Image
import gradio as gr
import transcriber
import json
import pandas as pd
import matplotlib.pyplot as plt
from scipy.io import wavfile
from scipy.signal import spectrogram
import numpy as np
from torch import nn

engine = transcriber.transcribe_SA(model_path='models/SA',verbose=0)
phonemizer = Phonemize.phonemization()
arpa2ipa = pd.read_csv('data/arpa2ipa.csv', sep='\\s+', header=None, names=['arpa','ipa'])

prompts = np.loadtxt('data/prompts.txt', dtype=str)

Attributes = engine.att_list
df_output = None

def select_prompt():
    return random.choice(prompts)

def phonemize_prompt(prompt, is_ipa=False):
    phonemes = phonemizer.cmu_phonemize(prompt)
    phonemes = [ph.lower() for ph in phonemes]
    if is_ipa:
        phonemes = [arpa2ipa[arpa2ipa.arpa==ph].ipa.values[0] for ph in phonemes]
    return ' '.join(phonemes)

def diff_fn():
    return [('H','+'),('E','-'),('N',None),('\n', None),('F','-'),('Fgo','-'),('M','+')]

def recognizeAudio(audio_file, attributes, is_ipa=False):
    global df_output
    if is_ipa:
        p2att_matrix = 'data/p2att_en_us-ipa.csv'
    else:
        p2att_matrix = 'data/p2att_en_us-arpa.csv'
    output = engine.transcribe(audio_file, attributes= 'all', phonological_matrix_file=p2att_matrix, human_readable=False)
    records = []
    d = json.loads(output)
    phonemes = d['Phoneme']['symbols']
    records.append(['Phoneme']+phonemes)
    for att in d['Attributes']:
        records.append([att['Name']]+att['Pattern'])
    df = pd.DataFrame.from_records(records)
    df.fillna('', inplace=True)
    df_output = df
    return df[df[0].isin(['Phoneme']+list(attributes))].to_html(header=False, index=False)

 
#Get error by matching the expected sequence with the recognized one and return the output in a format that can be visualized by the gradio HighlightedText box
def get_error(exp_list, rec_list):
    exp_list = list(exp_list)
    rec_list = list(rec_list)
    vocab = set(exp_list+rec_list)
    w2c = dict(zip(vocab,range(len(vocab))))
    
    exp_out = [[a,None] for a in exp_list]
    rec_out = [[a,None] for a in rec_list]  
    exp_enc = ''.join([chr(w2c[c]) for c in exp_list])
    rec_enc = ''.join([chr(w2c[c]) for c in rec_list])

    for op, exp_i, rec_i in editops(exp_enc, rec_enc):
        if op == 'replace':
            exp_out[exp_i][1] = 'S'
            rec_out[rec_i][1] = 'S'
        elif op == 'insert':
            rec_out[rec_i][1] = 'I'
        elif op == 'delete':
            exp_out[exp_i][1] = 'D'
    
    diff_list = [['Expected:\t', None]] + exp_out + [['\n',None]] + [['Recognized :\t', None]] + rec_out
    return diff_list


def scale_vector(vector, new_min, new_max):
    min_val = min(vector)
    max_val = max(vector)
    scaled_vector = []
    for val in vector:
        scaled_val = ((val - min_val) * (new_max - new_min) / (max_val - min_val)) + new_min
        scaled_vector.append(scaled_val)
    return scaled_vector



def create_spectrogram_with_att(wav_file, att_contour, att ):
    # Read the WAV file
    sampling_rate, data = wavfile.read(wav_file)

    # Calculate the spectrogram
    f, t, Sxx = spectrogram(data, fs=sampling_rate)
    fig, axs = plt.subplots(2, 1, figsize=(10, 10), sharex=True)

    # Plot the spectrogram
    axs[0].pcolormesh(t, f, 10 * np.log10(Sxx), shading='gouraud')  # Use grayscale colormap
    #plt.colorbar(label='Intensity (dB)')
    axs[0].set_ylabel('Frequency (Hz)')
    axs[0].set_xlabel('Time (s)')
    axs[0].set_title(f'Spectrogram with {att} Contour')
    axs[0].set_ylim(0, 8000)  # Adjust the frequency range if necessary

    ax_att = axs[0].twinx()
    # Plot the att contour
    x_points = att_contour.shape[0]
    time_att = np.arange(0,  x_points * 0.02, 0.02)[:x_points]  # Assuming pitch_contour is sampled every 20 ms
    ax_att.plot(time_att, att_contour, color='blue', label=f'{att} Contour')
    ax_att.set_ylim(0,1)
    ax_att.legend()

    # Plot the waveform
    time = np.arange(0, len(data)) / sampling_rate
    axs[1].plot(time, data, color='blue')
    axs[1].set_ylabel('Amplitude')
    axs[1].set_xlabel('Time (s)')
    axs[1].set_title('Waveform')

    #plt.show()
    return fig

def plot_contour(audio_file, att):
    indx_n = engine.processor.tokenizer.convert_tokens_to_ids([f'n_{att}'])[0]
    indx_p = engine.processor.tokenizer.convert_tokens_to_ids([f'p_{att}'])[0]
    index_all = [engine.processor.tokenizer.pad_token_id, indx_n, indx_p]
    prob = nn.functional.softmax(engine.logits.squeeze()[:,index_all], dim=-1)
    att_contour = prob[:,-1]
    fig = create_spectrogram_with_att(audio_file, att_contour, att)
    return fig


with gr.Blocks() as gui:
    with gr.Tab("Main"):
        prompt = gr.Textbox(label='Prompt', value=select_prompt)
        get_prompt = gr.Button("Get Prompt")
        get_prompt.click(fn=select_prompt, outputs=prompt)

        with gr.Row():
            with gr.Column(scale=3):
                prompt_phonemes = gr.Textbox(label="Expected Phonemes", interactive=False)
            with gr.Column(scale=1):
                is_ipa = gr.Checkbox(label="IPA")
        
        get_phoneme = gr.Button("Get Phonemes")
        get_phoneme.click(fn=phonemize_prompt, inputs=[prompt, is_ipa], outputs=prompt_phonemes)
        
        record_audio = gr.Audio(sources=["microphone","upload"], type="filepath")
        att_list = gr.Dropdown(label="Select Attributes", choices=sorted(Attributes), value=['vowel', 'voiced', 'consonant'] ,multiselect=True)
        process = gr.Button("Process Audio")
        
        recognition = gr.HTML(label='Output')
        
        process.click(fn=recognizeAudio, inputs=[record_audio,att_list, is_ipa], outputs=recognition)
        
    with gr.Tab("Assessment"):
        assess = gr.Button("Assessment")
        diff = []
        for i in range(len(Attributes)+1):
            diff.append(gr.HighlightedText(
                    combine_adjacent=False,
                    show_legend=True,
                    color_map={"S": "red", "I": "green", "D":"blue"}, visible=False))
            
        def get_assessment(prompt_phonemes):#, recognized_phonemes, recognized_attributes):
            outputs = [gr.HighlightedText(visible=False)]*(df_output.shape[0])
            outputs[0] = gr.HighlightedText(label=f"Phoneme Assessment",
                                              value=get_error(prompt_phonemes.split(), df_output.iloc[0].values[1:]),
                                              visible=True)
            i = 1
            for j,r in df_output.iloc[1:].iterrows():
                convert = lambda ph: '-' if f'n_{att}' in engine.p2att_map[ph] else '+'
                att = r.iloc[0]
                exp_att = [convert(ph) for ph in prompt_phonemes.split()]
                rec_att = r.iloc[1:].values
                if ''.join(exp_att) != ''.join(rec_att):
                    outputs[i] = gr.HighlightedText(label=f"{att} Assessment",
                                              value=get_error(exp_att, rec_att),
                                              visible=True)
                    i += 1
                
            return outputs

        assess.click(fn=get_assessment, inputs= [prompt_phonemes], outputs=diff)    
        

    with gr.Tab("Analysis"):
        selected_att = gr.Dropdown( sorted(Attributes), label="Select an Attribute to plot", value='voiced', interactive=True)
        do_plot = gr.Button('Plot')
        plot_block = gr.Plot(label='Spectrogram with Attribute Contour')
        do_plot.click(plot_contour, inputs=[record_audio,selected_att], outputs=plot_block)

gui.launch()