Update app.py

app.py

@@ -1,36 +1,34 @@
 import gradio as gr
-from transformers import pipeline, Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
-import torch
+from transformers import pipeline
 import pandas as pd
 import plotly.express as px
-import soundfile as sf
 
 # ------------------------------
 # Load pretrained models
 # ------------------------------
-# Text classifier
 text_classifier = pipeline(
     "text-classification",
     model="j-hartmann/emotion-english-distilroberta-base",
-    top_k=None  # returns all scores
+    return_all_scores=True
 )
 
-# Audio classifier (Wav2Vec2)
-audio_model_name = "Dpngtm/wav2vec2-emotion-recognition"
-audio_processor = Wav2Vec2Processor.from_pretrained(audio_model_name)
-audio_model = Wav2Vec2ForSequenceClassification.from_pretrained(audio_model_name)
+audio_classifier = pipeline(
+    "audio-classification",
+    model="superb/wav2vec2-base-superb-er"
+)
 
 # ------------------------------
-# Map emotion to emoji
+# Emotion to Emoji mapping
 # ------------------------------
 EMOJI_MAP = {
-    "anger": "😡",
-    "disgust": "🤢",
-    "fear": "😨",
-    "joy": "😄",
-    "neutral": "😐",
+    "joy": "😊",
     "sadness": "😢",
-    "surprise": "😲"
+    "anger": "😠",
+    "fear": "😨",
+    "love": "❤️",
+    "surprise": "😲",
+    "disgust": "🤢",
+    "neutral": "😐"
 }
 
 # ------------------------------
@@ -61,11 +59,11 @@ def fuse_predictions(text_preds=None, audio_preds=None, w_text=0.5, w_audio=0.5)
     return {"fused_label": best[0], "fused_score": round(best[1], 3), "all_scores": scores}
 
 # ------------------------------
-# Bar chart function
+# Create bar chart with emojis
 # ------------------------------
 def make_bar_chart(scores_dict, title="Emotion Scores"):
     df = pd.DataFrame({
-        "Emotion": list(scores_dict.keys()),
+        "Emotion": [f"{EMOJI_MAP.get(k, '')} {k}" for k in scores_dict.keys()],
         "Score": list(scores_dict.values())
     })
     fig = px.bar(df, x="Emotion", y="Score", text="Score",
@@ -76,63 +74,42 @@ def make_bar_chart(scores_dict, title="Emotion Scores"):
     return fig
 
 # ------------------------------
-# Audio prediction helper
-# ------------------------------
-def predict_audio(audio_file):
-    speech, sr = sf.read(audio_file)
-    inputs = audio_processor(speech, sampling_rate=sr, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        logits = audio_model(**inputs).logits
-    probs = torch.nn.functional.softmax(logits, dim=-1).squeeze().tolist()
-    labels = [audio_model.config.id2label[i] for i in range(len(probs))]
-    return [{"label": l, "score": s} for l, s in zip(labels, probs)]
-
-# ------------------------------
-# Gradio prediction function
+# Prediction function
 # ------------------------------
 def predict(text, audio, w_text, w_audio):
     text_preds, audio_preds = None, None
-
     if text:
-        text_preds = text_classifier(text)
+        text_preds = text_classifier(text)[0]
     if audio:
-        audio_preds = predict_audio(audio)
-
+        audio_preds = audio_classifier(audio)
     fused = fuse_predictions(text_preds, audio_preds, w_text, w_audio)
 
-    # Final emotion with animated emoji
-    label = fused['fused_label']
-    emoji = EMOJI_MAP.get(label, "❓")
-    final_emotion = f"### {label.upper()} {emoji} \nScore: {fused['fused_score']}"
-    animation = f"<div style='font-size:80px; animation: bounce 1s infinite;'>{emoji}</div>"
-
-    # Charts
+    # Bar charts
     charts = []
     if text_preds:
         charts.append(make_bar_chart({p['label']: p['score'] for p in text_preds}, "Text Emotion Scores"))
     if audio_preds:
         charts.append(make_bar_chart({p['label']: p['score'] for p in audio_preds}, "Audio Emotion Scores"))
-    charts.append(make_bar_chart(fused['all_scores'], "Fused Emotion Scores"))
+    charts.append(make_bar_chart(fused['all_scores'], f"Fused Emotion Scores\nPrediction: {EMOJI_MAP.get(fused['fused_label'], '')} {fused['fused_label']}"))
 
-    return final_emotion + animation, charts
+    return charts
 
 # ------------------------------
-# Build Gradio app
+# Build Gradio interface with emojis
 # ------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎭 Multimodal Emotion Classification (Text + Speech)")
+    gr.Markdown("## 🎭 Multimodal Emotion Classification (Text + Speech) 😎")
 
     with gr.Row():
         with gr.Column():
-            txt = gr.Textbox(label="Text input", placeholder="Type something emotional...")
-            aud = gr.Audio(type="filepath", label="Upload speech (wav/mp3)")
-            w1 = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Text weight")
-            w2 = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Audio weight")
-            btn = gr.Button("Predict")
+            txt = gr.Textbox(label="📝 Text input", placeholder="Type something emotional...")
+            aud = gr.Audio(type="filepath", label="🎤 Upload speech (wav/mp3)")
+            w1 = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="🔹 Text weight (w_text)")
+            w2 = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="🔹 Audio weight (w_audio)")
+            btn = gr.Button("✨ Predict")
         with gr.Column():
-            final_label = gr.HTML(label="Predicted Emotion")
             chart_output = gr.Plot(label="Emotion Scores")
 
-    btn.click(fn=predict, inputs=[txt, aud, w1, w2], outputs=[final_label, chart_output]*3)
+    btn.click(fn=predict, inputs=[txt, aud, w1, w2], outputs=[chart_output]*3)  # 3 charts: text, audio, fused
 
 demo.launch()