Update app.py

app.py

@@ -2,82 +2,87 @@ import gradio as gr
 import torch
 import torchaudio
 from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification
-from queue import Queue
-import threading
-import numpy as np
 
-# Check for device
+# Initialize device and model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Model setup
 model_name = "Hatman/audio-emotion-detection"
 feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
-model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name).to(device)
-
-# Real-time audio processing setup
-def preprocess_audio_chunk(audio_chunk, sampling_rate):
-    resampled_waveform = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=16000)(audio_chunk)
-    return {'speech': resampled_waveform.numpy().flatten(), 'sampling_rate': 16000}
-
-def inference_chunk(audio_chunk, sampling_rate):
-    example = preprocess_audio_chunk(audio_chunk, sampling_rate)
-    inputs = feature_extractor(example['speech'], sampling_rate=16000, return_tensors="pt", padding=True)
+model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
+
+# Define emotion labels
+EMOTION_LABELS = {
+    0: "angry",
+    1: "disgust",
+    2: "fear",
+    3: "happy",
+    4: "neutral",
+    5: "sad",
+    6: "surprise"
+}
+
+def preprocess_audio(audio):
+    """Preprocess audio file for model input"""
+    waveform, sampling_rate = torchaudio.load(audio)
+    resampled_waveform = torchaudio.transforms.Resample(
+        orig_freq=sampling_rate, 
+        new_freq=16000
+    )(waveform)
+    return {
+        'speech': resampled_waveform.numpy().flatten(), 
+        'sampling_rate': 16000
+    }
+
+def inference(audio):
+    """Full inference function returning emotion, logits, and predicted IDs"""
+    example = preprocess_audio(audio)
+    inputs = feature_extractor(
+        example['speech'], 
+        sampling_rate=16000, 
+        return_tensors="pt", 
+        padding=True
+    )
+    
+    # Move inputs to appropriate device
     inputs = {k: v.to(device) for k, v in inputs.items()}
+    
     with torch.no_grad():
-        logits = model(**inputs).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    emotion = model.config.id2label[predicted_ids.item()]
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_ids = torch.argmax(logits, dim=-1)
+        
+    predicted_emotion = EMOTION_LABELS[predicted_ids.item()]
+    return predicted_emotion, logits.tolist(), predicted_ids.tolist()
+
+def inference_label(audio):
+    """Simplified inference function returning only the emotion label"""
+    emotion, _, _ = inference(audio)
     return emotion
 
-# Queue for processing audio chunks
-audio_queue = Queue()
-results_queue = Queue()
-
-# Thread for processing audio in real-time
-def audio_processing_thread():
-    while True:
-        if not audio_queue.empty():
-            audio_chunk, sampling_rate = audio_queue.get()
-            emotion = inference_chunk(audio_chunk, sampling_rate)
-            results_queue.put(emotion)
-
-processing_thread = threading.Thread(target=audio_processing_thread, daemon=True)
-processing_thread.start()
-
-# Gradio interface for real-time streaming
-def real_time_inference_live(microphone_audio):
-    waveform = torch.tensor(microphone_audio["array"]).float()
-    sampling_rate = microphone_audio["sampling_rate"]
-
-    # Chunk size in samples (5 seconds chunks)
-    chunk_size = int(5 * sampling_rate)
-
-    # Process each chunk and collect live emotions
-    emotions = []
-    for start in range(0, len(waveform), chunk_size):
-        end = min(start + chunk_size, len(waveform))
-        audio_chunk = waveform[start:end]
-        if audio_chunk.size(0) > 0:
-            audio_queue.put((audio_chunk, sampling_rate))
-
-        # Retrieve results from the results queue
-        while not results_queue.empty():
-            emotion = results_queue.get()
-            emotions.append(emotion)
-
-    return "\n".join(emotions)
-
+# Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Live Emotion Detection from Audio")
-
-    audio_input = gr.Audio(streaming=True, label="Real-Time Audio Input", type="numpy")
-    emotion_output = gr.Textbox(label="Detected Emotions", lines=10)
-
-    def stream_audio_live(audio):
-        return real_time_inference_live(audio)
-
-    audio_input.stream(stream_audio_live, outputs=emotion_output)
-
-    gr.Markdown("This application processes audio in 5-second chunks and detects emotions in real-time.")
-
+    gr.Markdown("# Audio Emotion Detection")
+    
+    with gr.Tab("Quick Analysis"):
+        gr.Interface(
+            fn=inference_label,
+            inputs=gr.Audio(type="filepath"),
+            outputs=gr.Label(label="Detected Emotion"),
+            title="Audio Emotion Analysis",
+            description="Upload or record audio to detect the emotional content."
+        )
+    
+    with gr.Tab("Detailed Analysis"):
+        gr.Interface(
+            fn=inference,
+            inputs=gr.Audio(type="filepath"),
+            outputs=[
+                gr.Label(label="Detected Emotion"),
+                gr.JSON(label="Confidence Scores"),
+                gr.JSON(label="Internal IDs")
+            ],
+            title="Audio Emotion Analysis (Detailed)",
+            description="Get detailed analysis including confidence scores for each emotion."
+        )
+
+# Launch the app
 demo.launch(share=True)