Create app.py

app.py

@@ -0,0 +1,176 @@
+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)
+        predicted_index = torch.argmax(logits, dim=1).item()
+
+    return emotion_classes[predicted_index]
+
+
+# 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
+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)}"
+
+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="Predicted Emotion"),
+    title="Robust Speech Emotion Recognition (SupCon)",
+    description="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."
+)
+
+if __name__ == "__main__":
+    demo.launch()