Create app.py

app.py

@@ -0,0 +1,204 @@
+import spaces
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+import librosa
+import numpy as np
+import gradio as gr
+
+
+# Vars
+repo_name = "ipetrousov/ser_supcon"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+weights_dict = "supcon_stage_ft_acc60.pt"
+
+# Classes
+class EmotionFeatureExtractorBase(nn.Module):
+    """
+    Core CNN to extract spatial features from Mel-Spectrograms.
+    Outputs a flat 256-dimensional feature vector.
+    """
+
+    def __init__(self):
+        super(EmotionFeatureExtractorBase, self).__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten = nn.Flatten()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.global_pool(x)
+        x = self.flatten(x)
+        return x # Output shape: (batch_size, 256)
+
+
+class SupConStage2(nn.Module):
+    """ Model B (Stage 2): Frozen Encoder + Linear Classifier. """
+
+    def __init__(self, trained_encoder, num_classes=8):
+        super(SupConStage2, self).__init__()
+
+        # Pass trained encoder from Stage 1
+        self.encoder = trained_encoder
+
+        # Freeze encoder weights
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+
+        # Linear classifier
+        self.classifier = nn.Linear(256, num_classes)
+
+    def forward(self, x):
+        # Disable gradient tracking for the encoder to save memory/compute
+        with torch.no_grad():
+            features = self.encoder(x)
+
+        return self.classifier(features)
+
+
+# Function
+def preprocess_drive_audio(file_path, target_sample_rate=22050, n_mels=128, noise_std=0.0):
+    """
+    Loads audio and generates a normalized Mel-spectrogram using librosa.
+    Returns a PyTorch tensor shaped (1, 1, Mels, Time).
+    """
+
+    y, sr = librosa.load(file_path, sr=target_sample_rate, mono=True)
+
+    mel_spectrogram = librosa.feature.melspectrogram(
+        y=y,
+        sr=target_sample_rate,
+        n_fft=512,
+        hop_length=256,
+        n_mels=n_mels,
+        fmax=8000
+    )
+
+    mel_spectrogram_db = librosa.power_to_db(mel_spectrogram, ref=1.0)
+
+    mean = np.mean(mel_spectrogram_db)
+    std = np.std(mel_spectrogram_db)
+    mel_normalized = (mel_spectrogram_db - mean) / (std + 1e-6)
+
+    final_tensor = torch.tensor(mel_normalized, dtype=torch.float32)
+    final_tensor = final_tensor.unsqueeze(0).unsqueeze(0)
+
+    # Noise Injection Pipeline
+    if noise_std > 0.0:
+        noise_matrix = torch.randn_like(final_tensor)
+        final_tensor = final_tensor + (noise_matrix * noise_std)
+
+    return final_tensor
+
+
+def run_local_inference(model, file_path, device, noise_std=0.0):
+    """
+    Runs a forward pass
+    """
+
+    emotion_classes = ["Neutral", "Calm", "Happy", "Sad", "Angry", "Fearful", "Disgust", "Surprised"]
+
+    input_tensor = preprocess_drive_audio(file_path, noise_std=noise_std)
+    input_tensor = input_tensor.to(device)
+
+    model.eval()
+    with torch.no_grad():
+        logits = model(input_tensor)
+        
+        # Get the top prediction index
+        predicted_index = torch.argmax(logits, dim=1).item()
+        top_class = emotion_classes[predicted_index]
+        
+        # Apply softmax to convert logits to probabilities between 0 and 1
+        probabilities = torch.softmax(logits, dim=1).squeeze().tolist()
+
+    # Map the probabilities to their respective class names for Gradio
+    prob_dict = {emotion_classes[i]: probabilities[i] for i in range(len(emotion_classes))}
+    
+    return top_class, prob_dict
+
+
+# Load SOTA SupConFN
+blank_encoder = EmotionFeatureExtractorBase()
+model = SupConStage2(trained_encoder=blank_encoder, num_classes=8)
+
+# Download & Load weights
+model_path = hf_hub_download(repo_id=repo_name, filename=weights_dict)
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+model = model.to(device)
+
+
+# Gradio Interface Setup
+@spaces.GPU
+def predict_emotion(audio_path, noise_std):
+    if audio_path is None:
+        return "Upload an audio file."
+    
+    try:
+        pred = run_local_inference(model, audio_path, device, noise_std=noise_std)
+        return pred
+    except Exception as e:
+        return f"Error processing audio: {str(e)}"
+
+
+description_html = """
+Upload a short audio clip and test the architecture's resilience. The slider injects mathematically controlled Gaussian noise into the normalized Mel-spectrogram before classification.
+
+<div style="display: flex; gap: 10px; margin-top: 15px;">
+    <a href="https://github.com/gpetrousov/dl_assignment_demokritos" target="_blank">
+        <img src="https://img.shields.io/badge/GitHub-View_Repository-181717?style=for-the-badge&logo=github" alt="GitHub Repository" />
+    </a>
+    <a href="https://docs.google.com/presentation/d/1ShSmXY2OqCEbQQjYiLYCFc8Psu2bioECoFHhliiDEiQ/edit?usp=drive_link" target="_blank">
+        <img src="https://img.shields.io/badge/Presentation-View_Slides-F4B400?style=for-the-badge&logo=googleslides&logoColor=white" alt="Presentation Slides" />
+    </a>
+</div>
+"""
+
+demo = gr.Interface(
+    fn=predict_emotion,
+    inputs=[
+        gr.Audio(type="filepath", label="Upload Audio (.mp3 or .wav)"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label="Gaussian Noise Injection (STD)")
+    ],
+    outputs=[
+        gr.Textbox(label="Final Prediction"),
+        gr.Label(num_top_classes=8, label="Probability Distribution")
+    ],
+    title="Robust Speech Emotion Recognition (SupCon)",
+    description=description_html,
+)
+
+if __name__ == "__main__":
+    demo.launch()