app.py

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(ssr_mode=False)