Update app.py

app.py

@@ -3,58 +3,56 @@ import torch
 import torch.nn.functional as F
 from transformers import Wav2Vec2Processor, Wav2Vec2ForSequenceClassification
 import torchaudio
-import numpy as np
 
-# Define emotion labels
+# Define emotion labels and corresponding icons
 emotion_labels = ["angry", "calm", "disgust", "fearful", "happy", "neutral", "sad", "surprised"]
+emotion_icons = {
+    "angry": "😠", "calm": "😌", "disgust": "🤢", "fearful": "😨",
+    "happy": "😊", "neutral": "😐", "sad": "😢", "surprised": "😲"
+}
 
 # Load model and processor
 model_name = "Dpngtm/wav2vec2-emotion-recognition"
 model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
 processor = Wav2Vec2Processor.from_pretrained(model_name, num_labels=len(emotion_labels))
 
-# Define device
+# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 model.eval()
 
-# At the top with other global variables
-emotion_icons = {
-    "angry": "😠",
-    "calm": "😌",
-    "disgust": "🤢",
-    "fearful": "😨",
-    "happy": "😊",
-    "neutral": "😐",
-    "sad": "😢",
-    "surprised": "😲"
-}
-
 def recognize_emotion(audio):
     try:
+        # Handle case where no audio is provided
         if audio is None:
-            return {f"{emotion} {emotion_icons[emotion]}": 0 for emotion in emotion_labels}
-            
+            return {f"{emotion} {emotion_icons[emotion]}": 0.0 for emotion in emotion_labels}
+        
+        # Load and preprocess the audio
         audio_path = audio if isinstance(audio, str) else audio.name
         speech_array, sampling_rate = torchaudio.load(audio_path)
         
+        # Limit audio length to 1 minute (60 seconds)
         duration = speech_array.shape[1] / sampling_rate
         if duration > 60:
             return {
                 "Error": "Audio too long (max 1 minute)",
-                **{f"{emotion} {emotion_icons[emotion]}": 0 for emotion in emotion_labels}
+                **{f"{emotion} {emotion_icons[emotion]}": 0.0 for emotion in emotion_labels}
             }
         
+        # Resample audio if not at 16kHz
         if sampling_rate != 16000:
             resampler = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=16000)
             speech_array = resampler(speech_array)
         
+        # Convert stereo to mono if necessary
         if speech_array.shape[0] > 1:
             speech_array = torch.mean(speech_array, dim=0, keepdim=True)
-            
+        
+        # Normalize audio
         speech_array = speech_array / torch.max(torch.abs(speech_array))
         speech_array = speech_array.squeeze().numpy()
         
+        # Process audio with the model
         inputs = processor(speech_array, sampling_rate=16000, return_tensors='pt', padding=True)
         input_values = inputs.input_values.to(device)
         
@@ -62,32 +60,28 @@ def recognize_emotion(audio):
             outputs = model(input_values)
             logits = outputs.logits
             probs = F.softmax(logits, dim=-1)[0].cpu().numpy()
-
-            # Ensure probabilities sum to 1 and convert to percentages
-            probs = probs / probs.sum()  # Normalize to ensure sum is 1
             
+            # Convert probabilities to percentages and format results
             confidence_scores = {
-                f"{emotion} {emotion_icons[emotion]}": float(prob * 100)
+                f"{emotion} {emotion_icons[emotion]}": round(float(prob * 100), 2)
                 for emotion, prob in zip(emotion_labels, probs)
             }
             
-            sorted_scores = dict(sorted(
-                confidence_scores.items(), 
-                key=lambda x: x[1], 
-                reverse=True
-            ))
-            
+            # Sort scores in descending order
+            sorted_scores = dict(sorted(confidence_scores.items(), key=lambda x: x[1], reverse=True))
             return sorted_scores
-            
+    
     except Exception as e:
+        # Return error message along with zeroed-out emotion scores
         return {
             "Error": str(e),
-            **{f"{emotion} {emotion_icons[emotion]}": 0 for emotion in emotion_labels}
+            **{f"{emotion} {emotion_icons[emotion]}": 0.0 for emotion in emotion_labels}
         }
 
-# Create a formatted string of supported emotions
+# Supported emotions for display
 supported_emotions = " | ".join([f"{emotion_icons[emotion]} {emotion}" for emotion in emotion_labels])
 
+# Gradio Interface setup
 interface = gr.Interface(
     fn=recognize_emotion,
     inputs=gr.Audio(
@@ -115,11 +109,10 @@ interface = gr.Interface(
     """
 )
 
-
 if __name__ == "__main__":
     interface.launch(
         share=True,
         debug=True,
         server_name="0.0.0.0",
         server_port=7860
-    )
+    )