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.
GitHub Repository Presentation Slides
""" 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(ssr_mode=False)