Update app.py

app.py

@@ -1,16 +1,10 @@
 import gradio as gr
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import torchaudio
 from transformers import AutoConfig, Wav2Vec2Processor, Wav2Vec2FeatureExtractor
 from src.models import Wav2Vec2ForSpeechClassification
-
-import librosa
-import IPython.display as ipd
 import numpy as np
-import pandas as pd
-import os
 
 model_name_or_path = "andromeda01111/Malayalam_SA"
 config = AutoConfig.from_pretrained(model_name_or_path)
@@ -18,47 +12,39 @@ feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)
 sampling_rate = feature_extractor.sampling_rate
 model = Wav2Vec2ForSpeechClassification.from_pretrained(model_name_or_path)
 
-
 def speech_file_to_array_fn(path, sampling_rate):
     speech_array, _sampling_rate = torchaudio.load(path)
-    resampler = torchaudio.transforms.Resample(_sampling_rate)
+    resampler = torchaudio.transforms.Resample(_sampling_rate, sampling_rate)
     speech = resampler(speech_array).squeeze().numpy()
     return speech
 
-
-def predict(path, sampling_rate):
-    speech = speech_file_to_array_fn(path, sampling_rate)
+def predict(audio_path):
+    speech = speech_file_to_array_fn(audio_path, sampling_rate)
     features = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors="pt", padding=True)
-
+    
     input_values = features.input_values
     attention_mask = features.attention_mask
-
+    
     with torch.no_grad():
         logits = model(input_values, attention_mask=attention_mask).logits
-
+    
     scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0]
-    output_emotion = [{"Emotion": config.id2label[i], "Score": f"{round(score * 100, 3):.1f}%"} for i, score in enumerate(scores)]
-
+    output_emotion = {config.id2label[i]: f"{round(score * 100, 3):.1f}%" for i, score in enumerate(scores)}
+    
     return output_emotion
 
-
-# Wrapper function for Gradio
 def gradio_predict(audio):
-    predictions = predict(audio)
-    return [f"{pred['Emotion']}: {pred['Score']}" for pred in predictions]
-
-
-# Gradio interface
-emotions = [config.id2label[i] for i in range(len(config.id2label))]
-outputs = [gr.Textbox(label=emotion, interactive=False) for emotion in emotions]
+    return predict(audio)
 
+# Gradio Interface with Audio Recording (max duration: 10 seconds)
 interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Audio(label="Upload Audio", type="filepath"),
-    outputs=outputs,
+    fn=gradio_predict,
+    inputs=gr.Audio(source="microphone", type="filepath", label="Record or Upload Audio", streaming=False),
+    outputs=gr.JSON(label="Emotion Scores"),
     title="Emotion Recognition",
-    description="Upload an audio file to predict emotions and their corresponding percentages.",
+    description="Record or upload an audio file (max 10 sec) to predict emotions and their corresponding percentages.",
+    live=False,
 )
 
 # Launch the app
-interface.launch()
+interface.launch()