Update app.py

app.py

@@ -1,132 +1,168 @@
 import gradio as gr
 import numpy as np
-from textblob import TextBlob
-import speech_recognition as sr
+import torch
+from transformers import AutoModelForAudioClassification, AutoFeatureExtractor
+import librosa
+import os
+import warnings
+warnings.filterwarnings("ignore")
 
-class SentimentAnalyzer:
+class EmotionRecognizer:
     def __init__(self):
-        self.recognizer = sr.Recognizer()
-    
-    def audio_to_text(self, audio):
-        """Convert audio to text using speech recognition"""
-        try:
-            # Get audio data from Gradio input
-            sample_rate, audio_data = audio
-            
-            # Convert audio data to audio file format that speech_recognition can use
-            import io
-            import scipy.io.wavfile as wav
-            byte_io = io.BytesIO()
-            wav.write(byte_io, sample_rate, audio_data.astype(np.int16))
-            byte_io.seek(0)
-            
-            # Use speech recognition
-            with sr.AudioFile(byte_io) as source:
-                audio_data = self.recognizer.record(source)
-                text = self.recognizer.recognize_google(audio_data)
-                return text
-        except Exception as e:
-            return f"Error in speech recognition: {str(e)}"
-    
-    def analyze_sentiment(self, text):
-        """Analyze sentiment using TextBlob"""
+        # Initialize the model and feature extractor
+        self.model_name = "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition"
+        self.model = AutoModelForAudioClassification.from_pretrained(self.model_name)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        self.sample_rate = 16000
+        
+        # Define emotion labels
+        self.labels = ['angry', 'happy', 'sad', 'neutral', 'fearful']
+
+    def process_audio(self, audio):
+        """Process audio and return emotions with confidence scores"""
         try:
-            blob = TextBlob(text)
-            # Get polarity (-1 to 1) and subjectivity (0 to 1)
-            polarity = blob.sentiment.polarity
-            subjectivity = blob.sentiment.subjectivity
-            
-            # Determine sentiment category
-            if polarity > 0:
-                sentiment = "Positive"
-            elif polarity < 0:
-                sentiment = "Negative"
+            # Check if audio is a tuple (new Gradio audio format)
+            if isinstance(audio, tuple):
+                sample_rate, audio_data = audio
             else:
-                sentiment = "Neutral"
-            
-            # Format results
-            results_text = f"""
-Detected Text: "{text}"
-
-Analysis Results:
-- Overall Sentiment: {sentiment}
-- Polarity Score: {polarity:.2f} (-1 to +1)
-- Subjectivity Score: {subjectivity:.2f} (0 to 1)
-"""
-            
+                return "Error: Invalid audio format", None
+
+            # Resample if necessary
+            if sample_rate != self.sample_rate:
+                audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=self.sample_rate)
+
+            # Convert to float32 if not already
+            audio_data = audio_data.astype(np.float32)
+
+            # Extract features
+            inputs = self.feature_extractor(
+                audio_data,
+                sampling_rate=self.sample_rate,
+                return_tensors="pt",
+                padding=True
+            ).to(self.device)
+
+            # Get model predictions
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+
+            # Process results
+            scores = predictions[0].cpu().numpy()
+            results = [
+                {"label": label, "score": float(score)}
+                for label, score in zip(self.labels, scores)
+            ]
+
+            # Sort by confidence
+            results.sort(key=lambda x: x["score"], reverse=True)
+
+            # Format results for display
+            output_text = "Emotion Analysis Results:\n\n"
+            output_text += "\n".join([
+                f"{result['label'].title()}: {result['score']*100:.2f}%"
+                for result in results
+            ])
+
             # Prepare plot data
             plot_data = {
-                "labels": ["Polarity", "Subjectivity"],
-                "values": [polarity * 100, subjectivity * 100]  # Convert to percentage for visualization
+                "labels": [r["label"].title() for r in results],
+                "values": [r["score"] * 100 for r in results]
             }
-            
-            return results_text, plot_data
-            
+
+            return output_text, plot_data
+
         except Exception as e:
-            return f"Error in sentiment analysis: {str(e)}", None
+            return f"Error processing audio: {str(e)}", None
 
 def create_interface():
-    analyzer = SentimentAnalyzer()
+    # Initialize the emotion recognizer
+    recognizer = EmotionRecognizer()
     
-    def process_audio(audio):
+    # Define processing function for Gradio
+    def process_audio_file(audio):
         if audio is None:
             return "Please provide an audio input.", None
+            
+        output_text, plot_data = recognizer.process_audio(audio)
         
-        # Convert audio to text
-        text = analyzer.audio_to_text(audio)
-        if text.startswith("Error"):
-            return text, None
-        
-        # Analyze sentiment
-        return analyzer.analyze_sentiment(text)
+        if plot_data is not None:
+            return (
+                output_text,
+                gr.BarPlot.update(
+                    value=plot_data,
+                    x="labels",
+                    y="values",
+                    title="Emotion Confidence Scores",
+                    x_title="Emotions",
+                    y_title="Confidence (%)"
+                )
+            )
+        return output_text, None
 
-    # Create Gradio interface
-    with gr.Blocks() as interface:
-        gr.Markdown("# 🎤 Speech Sentiment Analysis")
+    # Create the Gradio interface
+    with gr.Blocks(title="Audio Emotion Recognition") as interface:
+        gr.Markdown("# 🎭 Audio Emotion Recognition")
         gr.Markdown("""
-        Speak or upload an audio file to analyze its emotional content.
-        The system will convert speech to text and analyze the sentiment.
+        Upload an audio file or record directly to analyze the emotional content.
+        The model will detect emotions like angry, happy, sad, neutral, and fearful.
         """)
         
         with gr.Row():
             with gr.Column():
+                # Input audio component (updated format)
                 audio_input = gr.Audio(
                     label="Upload or Record Audio",
                     type="numpy",
                     sources=["microphone", "upload"]
                 )
-                analyze_btn = gr.Button("Analyze Sentiment")
+                
+                # Process button
+                process_btn = gr.Button("Analyze Emotion", variant="primary")
             
             with gr.Column():
+                # Output components
                 output_text = gr.Textbox(
                     label="Analysis Results",
-                    lines=8
+                    lines=6
                 )
                 output_plot = gr.BarPlot(
-                    title="Sentiment Scores",
-                    x_title="Metrics",
-                    y_title="Score (%)"
+                    title="Emotion Confidence Scores",
+                    x_title="Emotions",
+                    y_title="Confidence (%)"
                 )
         
-        analyze_btn.click(
-            fn=process_audio,
+        # Set up event handler
+        process_btn.click(
+            fn=process_audio_file,
             inputs=[audio_input],
             outputs=[output_text, output_plot]
         )
         
         gr.Markdown("""
-        ### How to Use:
-        1. Click the microphone button to record or upload an audio file
-        2. Click "Analyze Sentiment" to process
-        3. View the results showing:
-           - Detected text from speech
-           - Overall sentiment (Positive/Negative/Neutral)
-           - Polarity score (-100% to +100%)
-           - Subjectivity score (0% to 100%)
+        ### Usage Instructions:
+        1. Click the microphone button to record audio or upload an audio file
+        2. Click "Analyze Emotion" to process the audio
+        3. View the results and confidence scores
+        
+        ### Notes:
+        - For best results, ensure clear audio with minimal background noise
+        - Speak naturally and clearly when recording
+        - The model works best with speech in English
         """)
     
     return interface
 
+def main():
+    # Create and launch the interface
+    interface = create_interface()
+    interface.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )
+
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch(share=True)
+    main()