Audio Classifier — Voice Sentiment Analysis

"""
Audio Classifier — Voice Sentiment Analysis
Trains a 2D CNN on Mel-spectrograms to classify speech emotions.

Dataset: RAVDESS (Ryerson Audio-Visual Database of Emotional Speech and Song)
Classes: neutral, happy, sad, angry, fear, surprise
"""

import os
import numpy as np
import librosa
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder

# ----- Configuration -----
SAMPLE_RATE = 22050
DURATION = 3.0          # seconds per clip
N_MFCC = 13
N_MELS = 128
HOP_LENGTH = 512
NUM_CLASSES = 6
BATCH_SIZE = 32
EPOCHS = 30
LR = 1e-3
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

EMOTIONS = ["neutral", "happy", "sad", "angry", "fear", "surprise"]


# ----- Feature extraction -----
def extract_features(file_path: str) -> np.ndarray:
    """Load audio file and extract a 47-D feature vector."""
    y, sr = librosa.load(file_path, sr=SAMPLE_RATE, duration=DURATION)

    # Pad or trim to fixed length
    target_len = int(SAMPLE_RATE * DURATION)
    if len(y) < target_len:
        y = np.pad(y, (0, target_len - len(y)))
    else:
        y = y[:target_len]

    # Pre-emphasis filter
    y = np.append(y[0], y[1:] - 0.97 * y[:-1])

    # MFCC
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=N_MFCC)
    mfcc_mean = mfcc.mean(axis=1)              # 13
    mfcc_std = mfcc.std(axis=1)                # 13

    # Chroma
    chroma = librosa.feature.chroma_stft(y=y, sr=sr).mean(axis=1)   # 12

    # Spectral features
    centroid = librosa.feature.spectral_centroid(y=y, sr=sr).mean()  # 1
    bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr).mean()
    rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr).mean()
    contrast = librosa.feature.spectral_contrast(y=y, sr=sr).mean()
    flatness = librosa.feature.spectral_flatness(y=y).mean()
    zcr = librosa.feature.zero_crossing_rate(y).mean()

    # Energy & pitch
    rms = librosa.feature.rms(y=y).mean()
    pitches, _ = librosa.piptrack(y=y, sr=sr)
    pitch = pitches[pitches > 0].mean() if (pitches > 0).any() else 0.0

    return np.concatenate([
        mfcc_mean, mfcc_std, chroma,
        [centroid, bandwidth, rolloff, contrast, flatness, zcr, rms, pitch]
    ]).astype(np.float32)


def to_melspectrogram(file_path: str) -> np.ndarray:
    """Convert audio to a normalized Mel-spectrogram (used as CNN input)."""
    y, sr = librosa.load(file_path, sr=SAMPLE_RATE, duration=DURATION)
    target_len = int(SAMPLE_RATE * DURATION)
    if len(y) < target_len:
        y = np.pad(y, (0, target_len - len(y)))
    else:
        y = y[:target_len]

    mel = librosa.feature.melspectrogram(
        y=y, sr=sr, n_mels=N_MELS, hop_length=HOP_LENGTH
    )
    mel_db = librosa.power_to_db(mel, ref=np.max)
    # Normalize to [0, 1]
    mel_db = (mel_db - mel_db.min()) / (mel_db.max() - mel_db.min() + 1e-8)
    return mel_db.astype(np.float32)


# ----- Dataset -----
class RAVDESSDataset(Dataset):
    """Dataset wrapper for Mel-spectrogram + emotion label."""

    def __init__(self, file_paths, labels):
        self.file_paths = file_paths
        self.labels = labels

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        mel = to_melspectrogram(self.file_paths[idx])
        # Add channel dimension: (1, n_mels, n_frames)
        mel = torch.from_numpy(mel).unsqueeze(0)
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        return mel, label


# ----- Model: 2D CNN -----
class AudioCNN(nn.Module):
    """Convolutional network for Mel-spectrogram classification."""

    def __init__(self, num_classes: int = NUM_CLASSES):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.3)

        # Compute flattened size once (depends on input dims)
        self._flat_size = self._get_flat_size()

        self.fc1 = nn.Linear(self._flat_size, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def _get_flat_size(self) -> int:
        with torch.no_grad():
            x = torch.zeros(1, 1, N_MELS, 130)  # sample input
            x = self.pool(F.relu(self.bn1(self.conv1(x))))
            x = self.pool(F.relu(self.bn2(self.conv2(x))))
            x = self.pool(F.relu(self.bn3(self.conv3(x))))
            return x.numel()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool(F.relu(self.bn1(self.conv1(x))))   # -> 32 × 64 × 65
        x = self.pool(F.relu(self.bn2(self.conv2(x))))   # -> 64 × 32 × 32
        x = self.pool(F.relu(self.bn3(self.conv3(x))))   # -> 128 × 16 × 16
        x = x.view(x.size(0), -1)
        x = self.dropout(F.relu(self.fc1(x)))
        return self.fc2(x)


# ----- Training loop -----
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    total_loss, correct = 0.0, 0
    for x, y in loader:
        x, y = x.to(DEVICE), y.to(DEVICE)
        optimizer.zero_grad()
        logits = model(x)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * x.size(0)
        correct += (logits.argmax(1) == y).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)


@torch.no_grad()
def evaluate(model, loader):
    model.eval()
    correct, total = 0, 0
    for x, y in loader:
        x, y = x.to(DEVICE), y.to(DEVICE)
        logits = model(x)
        correct += (logits.argmax(1) == y).sum().item()
        total += y.size(0)
    return correct / total


def main(data_dir: str):
    # Discover audio files (RAVDESS naming: 03-01-EMOTION-...)
    files, labels = [], []
    le = LabelEncoder()
    le.fit(EMOTIONS)

    for fname in sorted(os.listdir(data_dir)):
        if not fname.endswith(".wav"):
            continue
        # Example: "03-01-04-01-01-01-01.wav" -> emotion at index 2
        parts = fname.split("-")
        if len(parts) < 3:
            continue
        emotion_code = int(parts[2])
        # RAVDESS codes: 01=neutral, 03=happy, 04=sad, 05=angry, 06=fear, 08=surprise
        code_map = {1: 0, 3: 1, 4: 2, 5: 3, 6: 4, 8: 5}
        if emotion_code not in code_map:
            continue
        files.append(os.path.join(data_dir, fname))
        labels.append(code_map[emotion_code])

    print(f"Found {len(files)} files across {len(set(labels))} classes")

    # 80/20 split
    X_train, X_test, y_train, y_test = train_test_split(
        files, labels, test_size=0.2, random_state=42, stratify=labels
    )

    train_ds = RAVDESSDataset(X_train, y_train)
    test_ds = RAVDESSDataset(X_test, y_test)
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, num_workers=2)

    model = AudioCNN().to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
    criterion = nn.CrossEntropyLoss()

    best_acc = 0.0
    for epoch in range(1, EPOCHS + 1):
        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
        test_acc = evaluate(model, test_loader)
        scheduler.step()
        print(
            f"Epoch {epoch:02d}/{EPOCHS}  "
            f"loss={train_loss:.4f}  train_acc={train_acc:.3f}  test_acc={test_acc:.3f}"
        )
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(model.state_dict(), "audio_emotion_cnn.pt")

    print(f"\nBest test accuracy: {best_acc:.3f}")
    print("Saved model → audio_emotion_cnn.pt")


if __name__ == "__main__":
    main(data_dir="./data/ravdess")

/**
 * Real-time Audio Emotion Classifier
 * Captures microphone audio, extracts MFCC features, and predicts emotion
 * at ~30 FPS using TensorFlow.js.
 */

const EMOTIONS = ['neutral', 'happy', 'sad', 'angry', 'fear', 'surprise'];
const SAMPLE_RATE = 22050;
const FFT_SIZE = 2048;
const HOP_LENGTH = 512;
const N_MFCC = 13;
const N_MELS = 128;

class AudioEmotionClassifier {
  constructor() {
    this.audioContext = null;
    this.analyser = null;
    this.model = null;
    this.buffer = [];
    this.sampleRate = SAMPLE_RATE;
    this.isRunning = false;
    this.onPrediction = null;   // callback
  }

  /** Load the trained TensorFlow.js model. */
  async loadModel(modelPath = './audio_emotion_cnn/model.json') {
    this.model = await tf.loadLayersModel(modelPath);
    console.log('Model ready');
  }

  /** Start microphone capture and analysis loop. */
  async start() {
    if (this.isRunning) return;
    this.isRunning = true;

    // 1. Request microphone permission
    const stream = await navigator.mediaDevices.getUserMedia({
      audio: { echoCancellation: true, noiseSuppression: true }
    });

    // 2. Build audio graph
    this.audioContext = new AudioContext({ sampleRate: this.sampleRate });
    const source = this.audioContext.createMediaStreamSource(stream);
    this.analyser = this.audioContext.createAnalyser();
    this.analyser.fftSize = FFT_SIZE;
    this.analyser.smoothingTimeConstant = 0.6;
    source.connect(this.analyser);

    // 3. Start prediction loop
    this._loop();
  }

  stop() {
    this.isRunning = false;
    if (this.audioContext) this.audioContext.close();
  }

  /** Run feature extraction + prediction on every animation frame. */
  _loop() {
    if (!this.isRunning) return;

    const freqData = new Uint8Array(this.analyser.frequencyBinCount);
    const timeData = new Float32Array(this.analyser.fftSize);

    this.analyser.getByteFrequencyData(freqData);
    this.analyser.getFloatTimeDomainData(timeData);

    // Compute features and predict
    const features = this._extractFeatures(freqData, timeData);
    const prediction = this._predict(features);

    if (this.onPrediction) {
      this.onPrediction(prediction);
    }

    requestAnimationFrame(() => this._loop());
  }

  /**
   * Extract MFCC + spectral features from raw frequency bins.
   * Returns a Float32Array of length 47.
   */
  _extractFeatures(freqData, timeData) {
    // 1. Compute mel-spectrogram from FFT bins
    const melSpec = this._computeMelSpectrogram(freqData);

    // 2. Apply DCT to get MFCCs
    const mfcc = this._dct(melSpec).slice(0, N_MFCC);

    // 3. Compute mean & std of MFCCs across time
    const mfccMean = this._mean(mfcc);
    const mfccStd = this._std(mfcc, mfccMean);

    // 4. Chroma features
    const chroma = this._computeChroma(freqData);

    // 5. Spectral statistics
    const energy = this._rms(timeData);
    const zcr = this._zeroCrossingRate(timeData);
    const centroid = this._spectralCentroid(freqData);
    const flatness = this._spectralFlatness(freqData);

    // 6. Pitch (fundamental frequency)
    const pitch = this._estimatePitch(timeData);

    // Concatenate into single feature vector
    return new Float32Array([
      ...mfccMean, ...mfccStd, ...chroma,
      centroid, zcr, flatness, energy, pitch
    ]);
  }

  _computeMelSpectrogram(freqData) {
    // Convert magnitude spectrum to mel scale using triangular filters
    const melFilters = this._buildMelFilterbank(N_MELS, freqData.length);
    const mel = new Float32Array(N_MELS);
    for (let m = 0; m < N_MELS; m++) {
      let sum = 0;
      for (let k = 0; k < freqData.length; k++) {
        sum += freqData[k] * melFilters[m * freqData.length + k];
      }
      mel[m] = Math.log(sum + 1e-8);   // log-mel
    }
    return mel;
  }

  _dct(mel) {
    // Type-II DCT, orthonormal
    const N = mel.length;
    const out = new Float32Array(N);
    for (let k = 0; k < N; k++) {
      let sum = 0;
      for (let n = 0; n < N; n++) {
        sum += mel[n] * Math.cos((Math.PI / N) * (n + 0.5) * k);
      }
      out[k] = sum * Math.sqrt(2 / N);
    }
    return out;
  }

  _mean(arr) {
    let s = 0;
    for (let i = 0; i < arr.length; i++) s += arr[i];
    return s / arr.length;
  }

  _std(arr, mean) {
    let v = 0;
    for (let i = 0; i < arr.length; i++) v += (arr[i] - mean) ** 2;
    return Math.sqrt(v / arr.length);
  }

  _zeroCrossingRate(signal) {
    let zc = 0;
    for (let i = 1; i < signal.length; i++) {
      if ((signal[i - 1] >= 0) !== (signal[i] >= 0)) zc++;
    }
    return zc / signal.length;
  }

  _rms(signal) {
    let s = 0;
    for (let i = 0; i < signal.length; i++) s += signal[i] ** 2;
    return Math.sqrt(s / signal.length);
  }

  _spectralCentroid(freqData) {
    let num = 0, den = 0;
    for (let k = 0; k < freqData.length; k++) {
      num += k * freqData[k];
      den += freqData[k];
    }
    return den > 0 ? num / den : 0;
  }

  _spectralFlatness(freqData) {
    let logSum = 0, count = 0;
    for (let k = 0; k < freqData.length; k++) {
      if (freqData[k] > 0) {
        logSum += Math.log(freqData[k] + 1);
        count++;
      }
    }
    const geoMean = count > 0 ? Math.exp(logSum / count) : 0;
    let arithSum = 0;
    for (let k = 0; k < freqData.length; k++) arithSum += freqData[k];
    const arithMean = arithSum / freqData.length;
    return arithMean > 0 ? geoMean / arithMean : 0;
  }

  /** Estimate pitch via autocorrelation peak detection. */
  _estimatePitch(signal) {
    const minLag = Math.floor(this.sampleRate / 500);   // 500 Hz max
    const maxLag = Math.floor(this.sampleRate / 80);    // 80 Hz min
    let bestLag = -1, bestCorr = 0;
    for (let lag = minLag; lag <= maxLag; lag++) {
      let corr = 0;
      for (let i = 0; i < signal.length - lag; i++) {
        corr += signal[i] * signal[i + lag];
      }
      if (corr > bestCorr) { bestCorr = corr; bestLag = lag; }
    }
    return bestLag > 0 ? this.sampleRate / bestLag : 0;
  }

  /** Run model inference and return per-class probabilities. */
  async _predict(features) {
    if (!this.model) return null;
    const input = tf.tensor2d([features]);
    const logits = this.model.predict(input);
    const probs = await logits.data();
    input.dispose();
    logits.dispose();
    return Array.from(probs);
  }
}

// ----- Usage Example -----
async function main() {
  const classifier = new AudioEmotionClassifier();
  await classifier.loadModel();

  classifier.onPrediction = (probs) => {
    if (!probs) return;
    const idx = probs.indexOf(Math.max(...probs));
    console.clear();
    console.log(`Detected: ${EMOTIONS[idx]} (${(probs[idx]*100).toFixed(1)}%)`);
    probs.forEach((p, i) =>
      console.log(`  ${EMOTIONS[i].padEnd(10)} ${(p*100).toFixed(1)}%`)
    );
  };

  await classifier.start();
}

main().catch(console.error);

/**
 * Batch inference pipeline for classifying audio files (WAV/MP3/OGG).
 * Loads a model, processes files, and returns ranked predictions.
 */

const fs = require('fs');
const path = require('path');
const { PCMDecoder } = require('pcm-decoder');
const ort = require('onnxruntime-node');   // ONNX Runtime

const EMOTIONS = ['neutral', 'happy', 'sad', 'angry', 'fear', 'surprise'];
const SAMPLE_RATE = 22050;
const DURATION = 3.0;
const N_MELS = 128;

/** Decode any audio file to a mono Float32Array at SAMPLE_RATE. */
async function decodeAudio(filePath) {
  const buffer = fs.readFileSync(filePath);
  const audioCtx = new AudioContext({ sampleRate: SAMPLE_RATE });
  const audioBuffer = await audioCtx.decodeAudioData(buffer);
  // Mix to mono
  const channelData = audioBuffer.getChannelData(0);
  // Pad or trim
  const target = SAMPLE_RATE * DURATION;
  const y = new Float32Array(target);
  for (let i = 0; i < Math.min(channelData.length, target); i++) {
    y[i] = channelData[i];
  }
  await audioCtx.close();
  return y;
}

/** Compute log-Mel-spectrogram for CNN input. */
function computeLogMel(y) {
  // FFT frame of 1024 samples, hop 512
  const N_FFT = 1024;
  const HOP = 512;
  const numFrames = Math.floor((y.length - N_FFT) / HOP) + 1;
  const mel = new Float32Array(N_MELS * numFrames);

  // Hann window
  const window = new Float32Array(N_FFT);
  for (let i = 0; i < N_FFT; i++) {
    window[i] = 0.5 * (1 - Math.cos((2 * Math.PI * i) / (N_FFT - 1)));
  }

  // Mel filterbank (simplified triangular)
  const melBasis = build_mel_filterbank(N_MELS, N_FFT / 2 + 1, SAMPLE_RATE);

  for (let f = 0; f < numFrames; f++) {
    // 1. Windowed signal
    const re = new Float32Array(N_FFT);
    const im = new Float32Array(N_FFT);
    for (let i = 0; i < N_FFT; i++) {
      re[i] = y[f * HOP + i] * window[i];
    }
    // 2. FFT (use a library like fft.js — omitted for brevity)
    fft(re, im);
    // 3. Power spectrum
    const power = new Float32Array(N_FFT / 2 + 1);
    for (let k = 0; k < power.length; k++) {
      power[k] = (re[k] ** 2 + im[k] ** 2);
    }
    // 4. Apply mel filterbank
    for (let m = 0; m < N_MELS; m++) {
      let s = 0;
      for (let k = 0; k < power.length; k++) {
        s += power[k] * melBasis[m * power.length + k];
      }
      mel[m * numFrames + f] = Math.log(s + 1e-8);
    }
  }
  // Normalize
  const min = Math.min(...mel), max = Math.max(...mel);
  for (let i = 0; i < mel.length; i++) {
    mel[i] = (mel[i] - min) / (max - min + 1e-8);
  }
  return { data: mel, frames: numFrames };
}

/** Run model and decode predictions. */
async function classify(filePath, session) {
  const y = await decodeAudio(filePath);
  const mel = computeLogMel(y);

  // Build ONNX tensor (1, 1, n_mels, n_frames)
  const tensor = new ort.Tensor('float32', mel.data, [1, 1, N_MELS, mel.frames]);
  const feeds = { input: tensor };
  const outputs = await session.run(feeds);

  const logitsArr = outputs.logits.data;
  const probs = softmax(Array.from(logitsArr));
  const ranked = probs
    .map((p, i) => ({ emotion: EMOTIONS[i], prob: p }))
    .sort((a, b) => b.prob - a.prob);

  return {
    file: path.basename(filePath),
    top: ranked[0],
    all: ranked,
  };
}

function softmax(arr) {
  const m = Math.max(...arr);
  const exps = arr.map(v => Math.exp(v - m));
  const s = exps.reduce((a, b) => a + b, 0);
  return exps.map(e => e / s);
}

async function main() {
  const session = await ort.InferenceSession.create('./audio_emotion_cnn.onnx');

  const files = fs.readdirSync('./audio_samples')
    .filter(f => /\.(wav|mp3|ogg)$/.test(f));

  console.log(`Classifying ${files.length} audio files...\n`);
  for (const f of files) {
    const result = await classify(`./audio_samples/${f}`, session);
    console.log(`${result.file.padEnd(30)} → ${result.top.emotion} (${(result.top.prob*100).toFixed(1)}%)`);
  }
}

main().catch(console.error);