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


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):
    #print(','.join(attributes))
    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= tuple(attributes), phonological_matrix_file=p2att_matrix, human_readable=False)
    records = []
    d = json.loads(output)
    phonemes = d['Phoneme']['symbols']
    #if is_ipa:
    #    phonemes = [arpa2ipa[arpa2ipa.arpa==ph].ipa.values[0] for ph in phonemes] 
    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.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, ax = plt.subplots(figsize=(10, 5))

    # Plot the spectrogram
    ax.pcolormesh(t, f, 10 * np.log10(Sxx), shading='gouraud')
    ax.set_ylabel('Frequency (Hz)')
    ax.set_xlabel('Time (s)')
    ax.set_title(f'Spectrogram with {att} Contour')
    ax.set_ylim(0, 8000)  # Adjust the frequency range if necessary

    # Plot the att contour
    time_pitch = np.arange(0, len(att_contour) * 0.02, 0.02)  # Assuming pitch_contour is sampled every 20 ms
    ax.plot(time_pitch, att_contour, color='blue', label=f'{att} Contour')
    ax.legend()

    return fig

def plot_contour(audio_file, att):
    indx = engine.processor.tokenizer.convert_tokens_to_ids([f'p_{att}'])
    att_contour = engine.logits.squeeze()[:,indx]
    att_contour = scale_vector(att_contour, 0, 6000)
    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=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)]*(len(Attributes)+1)
            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 i,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
                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( 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()