import numpy as np import pandas as pd from main import SpeechClassifierOutput, Wav2Vec2ForSpeechClassification from datasets import load_dataset from transformers import AutoConfig, Wav2Vec2Processor import torchaudio import torch import torch.nn.functional as F import seaborn as sns import matplotlib.pyplot as plt import streamlit as st import os sns.set_theme(style="darkgrid", palette="pastel") def demo_speech_file_to_array_fn(path): speech_array, _sampling_rate = torchaudio.load(path, normalize=True) resampler = torchaudio.transforms.Resample(_sampling_rate, 16_000) speech = resampler(speech_array).squeeze().numpy() return speech def demo_predict(df_row): path, emotion = df_row["path"], df_row["emotion"] speech = demo_speech_file_to_array_fn(path) features = processor(speech, sampling_rate=16_000, return_tensors="pt", padding=True) input_values = features.input_values.to(device) attention_mask = features.attention_mask.to(device) with torch.no_grad(): logits = model(input_values, attention_mask=attention_mask).logits scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0] outputs = [{"Emotion": config.id2label[i], "Score": round(score * 100, 3)} for i, score in enumerate(scores)] return outputs def cache_model(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_name_or_path = 'khizon/greek-speech-emotion-classifier-demo' generic_greek_model = 'lighteternal/wav2vec2-large-xlsr-53-greek' config = AutoConfig.from_pretrained(model_name_or_path) processor = Wav2Vec2Processor.from_pretrained(generic_greek_model) model = Wav2Vec2ForSpeechClassification.from_pretrained(model_name_or_path).to(device) return config, processor, model, device @st.cache def load_data(): return pd.read_csv('data/test.csv', delimiter = '\t') def bar_plot(df): fig = plt.figure(figsize=(8, 6)) plt.title("Prediction Scores") plt.xticks(fontsize=12) sns.barplot(x="Score", y="Emotion", data=df) st.pyplot(fig) if __name__ == '__main__': if not os.path.exists('/home/user/app/aesdd.zip'): os.system('python download_dataset.py') test = load_data() config, processor, model, device = cache_model() print('Model loaded') st.title("Emotion Classifier for Greek Speech Audio Demo") if st.button("Classify Random Audio"): # Load demo file idx = np.random.randint(0, len(test)) sample = test.iloc[idx] audio_file = open(sample['path'], 'rb') audio_bytes = audio_file.read() st.success(f'Label: {sample["emotion"]}') st.audio(audio_bytes, format='audio/ogg') outputs = demo_predict(sample) r = pd.DataFrame(outputs) # st.dataframe(r) bar_plot(r)