server.py

"""
SoundSense – Audio Anomaly Detection Server
Replicates the full pipeline from main.ipynb as a FastAPI service.
"""

import os
import sys
import json
import uuid
import shutil
import logging
import traceback
from pathlib import Path

# Workaround: transformers 4.57+ requires torch>=2.6 for safe deserialization.
# All models here are local/trusted, so we allow the older torch.load path.
os.environ["TRANSFORMERS_ALLOW_UNSAFE_DESERIALIZATION"] = "1"


import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
import joblib
import librosa
import soundfile as sf
import re
import nltk

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("soundsense")

BASE_DIR = Path(__file__).resolve().parent
UPLOAD_DIR = BASE_DIR / "uploads"
UPLOAD_DIR.mkdir(exist_ok=True)

SAMPLE_RATE = 16000
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# NLTK setup
nltk.data.path.append(str(BASE_DIR / 'assets' / 'nltk_data'))
nltk.download('stopwords', download_dir=str(BASE_DIR / 'assets' / 'nltk_data'), quiet=True)
nltk.download('punkt', download_dir=str(BASE_DIR / 'assets' / 'nltk_data'), quiet=True)
nltk.download('punkt_tab', download_dir=str(BASE_DIR / 'assets' / 'nltk_data'), quiet=True)
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize

stop_words = set(stopwords.words('english'))

def clean_text(text):
    text = re.sub(r'<.*?>', '', text)  # Remove HTML tags
    text = re.sub(r'[^a-zA-Z0-9\s]', '', text)  # Remove non-alphanumeric characters
    text = re.sub(r'http\S+|www\S+|https\S+', '', text, flags=re.MULTILINE)  # Remove URLs
    text = re.sub(r'\@\w+|\#', '', text)  # Remove mentions and hashtags
    text = re.sub(r'\&\w+;', '', text)  # Remove HTML entities
    text = re.sub(r'\d+', '', text)  # Remove digits
    text = re.sub(r'[^\w\s]', '', text)  # Remove punctuation
    text = text.lower()  # Convert to lowercase
    return text

def remove_stop_words(text):
    word_tokens = word_tokenize(text)
    filtered_text = [word for word in word_tokens if word not in stop_words]
    return ' '.join(filtered_text)

# Emotion label maps
text_emotions_map = {0: "sadness", 1: "joy", 2: "love", 3: "anger", 4: "fear", 5: "surprise"}
text_to_audio_map = {
    "sadness": ["SAD"],
    "joy": ["HAP"],
    "love": ["HAP", "NEU"],
    "anger": ["ANG", "DIS"],
    "fear": ["FEA", "DIS"],
    "surprise": ["HAP", "NEU"],
}
audio_emotion_map = {"HAP": 0, "SAD": 1, "ANG": 2, "FEA": 3, "DIS": 4, "NEU": 5}
wav2vec2_emotion_map = {"ANG": 0, "DIS": 1, "FEA": 2, "HAP": 3, "NEU": 4, "SAD": 5}
label_to_emotion = {0: "HAP", 1: "SAD", 2: "ANG", 3: "FEA", 4: "DIS", 5: "NEU"}
normalized_label_map = {"ANG": "angry", "DIS": "disgust", "FEA": "fear", "HAP": "happy", "NEU": "neutral", "SAD": "sad"}
audio_normalizer = {"angry": "ANG", "happy": "HAP", "sad": "SAD", "neutral": "NEU", "fear": "FEA", "disgust": "DIS"}

# ---------------------------------------------------------------------------
# Model loading (graceful – missing models are skipped)
# ---------------------------------------------------------------------------
from manage_models import AudioCNN, extract_features
from manage_models_v2 import convert_wav_audio, ml_model_extract_features
from torch import nn
from transformers import Wav2Vec2Model

# Transformers 4.57+ hard-blocks torch.load when torch < 2.6 (CVE-2025-32434).
# All models here are local/trusted, so we bypass the check instead of upgrading torch.
_noop = lambda: None
import transformers.utils.import_utils as _tu
if hasattr(_tu, "check_torch_load_is_safe"):
    _tu.check_torch_load_is_safe = _noop
import transformers.modeling_utils as _mu
if hasattr(_mu, "check_torch_load_is_safe"):
    _mu.check_torch_load_is_safe = _noop

# Helper: build the emotion recognition model with safetensors support
class EmotionRecognitionModelSafe(nn.Module):
    """Same architecture as EmotionRecognitionModel but loads wav2vec2-base via safetensors."""
    def __init__(self, num_labels):
        super().__init__()
        self.wav2vec2 = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base")
        self.classifier = nn.Sequential(
            nn.Linear(self.wav2vec2.config.hidden_size, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_labels),
        )

    def forward(self, input_values):
        outputs = self.wav2vec2(input_values)
        hidden_states = outputs.last_hidden_state
        pooled_output = hidden_states.mean(dim=1)
        return self.classifier(pooled_output)

# 1) ASR – Wav2Vec2 for transcription
logger.info("Loading ASR model (wav2vec2-base-960h)...")
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
asr_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
asr_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
asr_model.eval()
logger.info("ASR model loaded.")

# 2) RoBERTa text emotion
roberta_model_obj = None
roberta_tokenizer_obj = None
roberta_model_path = BASE_DIR / "roberta" / "model"
roberta_tokenizer_path = BASE_DIR / "roberta" / "tokenizer"
if roberta_model_path.exists() and roberta_tokenizer_path.exists():
    logger.info("Loading RoBERTa text-emotion model...")
    from transformers import RobertaTokenizer, RobertaForSequenceClassification
    roberta_model_obj = RobertaForSequenceClassification.from_pretrained(str(roberta_model_path))
    roberta_tokenizer_obj = RobertaTokenizer.from_pretrained(str(roberta_tokenizer_path))
    roberta_model_obj.eval()
    roberta_model_obj.to(device)
    logger.info("RoBERTa model loaded.")
else:
    logger.warning("RoBERTa model NOT found – text emotion prediction will be skipped.")

# 3) Wav2Vec2 audio emotion
wav2vec2_emotion_model = None
wav2vec2_model_path = BASE_DIR / "Wav2Vec2" / "best_model_v1.pt"
if wav2vec2_model_path.exists():
    logger.info("Loading Wav2Vec2 audio-emotion model...")
    from transformers import Wav2Vec2FeatureExtractor
    wav2vec2_emotion_model = EmotionRecognitionModelSafe(num_labels=len(audio_emotion_map))
    wav2vec2_emotion_model.load_state_dict(torch.load(str(wav2vec2_model_path), map_location=device, weights_only=False))
    wav2vec2_emotion_model.to(device)
    wav2vec2_emotion_model.eval()
    wav2vec2_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base")
    logger.info("Wav2Vec2 emotion model loaded.")
else:
    logger.warning("Wav2Vec2 emotion model NOT found – skipping.")

# 4) CNN audio emotion
cnn_model_obj = None
cnn_model_path = BASE_DIR / "CNN" / "cnn_model.pth"
if cnn_model_path.exists():
    logger.info("Loading CNN audio-emotion model...")
    cnn_model_obj = AudioCNN()
    cnn_model_obj.load_state_dict(torch.load(str(cnn_model_path), map_location=device, weights_only=False))
    cnn_model_obj.eval()
    cnn_model_obj.to(device)
    logger.info("CNN model loaded.")
else:
    logger.warning("CNN model NOT found – skipping.")

# 5) SVM
svm_model_obj = None
svm_model_path = BASE_DIR / "SVM" / "svm_model.joblib"
if svm_model_path.exists():
    logger.info("Loading SVM model...")
    svm_model_obj = joblib.load(str(svm_model_path))
    logger.info("SVM model loaded.")
else:
    logger.warning("SVM model NOT found – skipping.")

# 6) Random Forest
rf_model_obj = None
rf_model_path = BASE_DIR / "RF" / "rf_model.joblib"
if rf_model_path.exists():
    logger.info("Loading RF model...")
    rf_model_obj = joblib.load(str(rf_model_path))
    logger.info("RF model loaded.")
else:
    logger.warning("RF model NOT found – skipping.")

# Label encoder (shared by SVM & RF)
le = None
le_path = BASE_DIR / "SVM" / "label_encoder.joblib"
if le_path.exists():
    le = joblib.load(str(le_path))

# StandardScaler (shared by SVM & RF)
scaler = None
scaler_path = BASE_DIR / "SVM" / "scaler.joblib"
if scaler_path.exists():
    scaler = joblib.load(str(scaler_path))

# ---------------------------------------------------------------------------
# Pipeline functions
# ---------------------------------------------------------------------------

def transcribe_audio(file_path: str) -> str:
    """Transcribe audio to text using Wav2Vec2 ASR."""
    waveform, sample_rate = torchaudio.load(file_path)
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resampler(waveform)
    waveform = waveform.squeeze()
    inputs = asr_processor(waveform, sampling_rate=16000, return_tensors="pt", padding=True)
    with torch.no_grad():
        logits = asr_model(inputs.input_values).logits
    predicted_ids = torch.argmax(logits, dim=-1)
    return asr_processor.decode(predicted_ids[0]).lower()


def predict_text_emotion(text: str):
    """Predict emotion from text using RoBERTa."""
    if roberta_model_obj is None:
        return None, 0.0
    
    # Preprocess text exactly like training
    text = clean_text(text)
    text = remove_stop_words(text)
    
    inputs = roberta_tokenizer_obj(text, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = roberta_model_obj(**inputs)
    logits = outputs.logits
    prediction = logits.argmax(dim=-1)
    label = prediction.item()
    probs = F.softmax(logits, dim=-1)
    prob = probs[0][label].item()
    return text_emotions_map[label], prob


def predict_wav2vec2_emotion(file_path: str):
    """Predict audio emotion using Wav2Vec2."""
    if wav2vec2_emotion_model is None:
        return None
    waveform, sr = torchaudio.load(file_path)
    if sr != SAMPLE_RATE:
        waveform = torchaudio.transforms.Resample(orig_freq=sr, new_freq=SAMPLE_RATE)(waveform)
    waveform = waveform.squeeze(0)
    inputs = wav2vec2_feature_extractor(
        waveform.numpy(), sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True
    )
    input_values = inputs.input_values.to(device)
    with torch.no_grad():
        logits = wav2vec2_emotion_model(input_values)
        predicted_class_id = torch.argmax(logits, dim=-1).item()
    reverse_map = {v: k for k, v in wav2vec2_emotion_map.items()}
    return normalized_label_map[reverse_map[predicted_class_id]]


def predict_cnn_emotion(file_path: str):
    """Predict audio emotion using CNN."""
    if cnn_model_obj is None:
        return None
    mel = extract_features(file_path, augment=False)
    if mel is None:
        return None
    mel = torch.tensor(mel, dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
    mel = mel.to(device).permute(0, 3, 1, 2)
    with torch.no_grad():
        output = cnn_model_obj(mel)
        _, predicted = torch.max(output, 1)
    return normalized_label_map[label_to_emotion[predicted.item()]]


def predict_svm_emotion(file_path: str):
    """Predict audio emotion using SVM."""
    if svm_model_obj is None or le is None or scaler is None:
        return None
    features = ml_model_extract_features(file_path)
    if features is None:
        return None
    features = np.expand_dims(features, axis=0)
    features = scaler.transform(features)
    prediction = svm_model_obj.predict(features)
    return le.inverse_transform(prediction)[0]


def predict_rf_emotion(file_path: str):
    """Predict audio emotion using Random Forest."""
    if rf_model_obj is None or le is None or scaler is None:
        return None
    features = ml_model_extract_features(file_path)
    if features is None:
        return None
    features = np.expand_dims(features, axis=0)
    features = scaler.transform(features)
    prediction = rf_model_obj.predict(features)
    return le.inverse_transform(prediction)[0]


def evaluate_mismatch(result_map: dict) -> dict:
    """Run anomaly detection logic comparing text vs audio emotions."""
    text_emotion = (result_map.get("roberta_text_emotion") or "").lower()
    if not text_emotion:
        return {
            "status": "⚠️ Text emotion model unavailable – cannot run anomaly check.",
            "match_ratio": None,
            **result_map,
        }

    valid_audio_emotions = text_to_audio_map.get(text_emotion, [])

    audio_preds_raw = [
        result_map.get("wav2vec2_audio_emotion"),
        result_map.get("cnn_audio_emotion"),
        result_map.get("svm_audio_emotion"),
        result_map.get("rf_audio_emotion"),
    ]
    audio_preds = [p for p in audio_preds_raw if p is not None]

    if not audio_preds:
        return {
            "status": "⚠️ No audio emotion models available – cannot run anomaly check.",
            "match_ratio": None,
            **result_map,
        }

    normalized_preds = [audio_normalizer.get(pred.lower()) for pred in audio_preds if pred]
    matches = [p for p in normalized_preds if p in valid_audio_emotions]
    match_ratio = len(matches) / len(normalized_preds) if normalized_preds else 0

    if match_ratio == 1.0:
        status = "✅ No vocal-text mismatch detected. Emotion is consistent across modalities."
    elif match_ratio >= 0.5:
        status = "⚠️ Partial mismatch: Some models disagree. Possible nuanced expression or mild anomaly."
    else:
        status = "❗ Mismatch detected: Text and voice emotions do not align clearly. Possible vocal anomaly."

    return {
        "status": status,
        "match_ratio": match_ratio,
        "text_emotion": text_emotion,
        "expected_audio_emotions": valid_audio_emotions,
        "audio_predictions": normalized_preds,
        **result_map,
    }


def normalize_audio(file_path: str) -> str:
    """Normalize audio (mono, 16 kHz) and return the path to the normalised file."""
    output_dir = str(UPLOAD_DIR / "normalised")
    _, _, output_path = convert_wav_audio(file_path, output_dir=output_dir)
    return output_path


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------
app = FastAPI(title="SoundSense – Audio Anomaly Detection")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/", response_class=HTMLResponse)
async def index():
    html_path = BASE_DIR / "frontend" / "index.html"
    return HTMLResponse(content=html_path.read_text(encoding="utf-8"))


@app.get("/health")
async def health():
    models = {
        "asr": True,
        "roberta_text_emotion": roberta_model_obj is not None,
        "wav2vec2_audio_emotion": wav2vec2_emotion_model is not None,
        "cnn_audio_emotion": cnn_model_obj is not None,
        "svm_audio_emotion": svm_model_obj is not None,
        "rf_audio_emotion": rf_model_obj is not None,
    }
    return {"status": "ok", "device": str(device), "models_loaded": models}


@app.post("/analyze")
async def analyze(file: UploadFile = File(...)):
    """Upload a WAV audio file and get the full anomaly analysis."""
    # Validate file type
    if not file.filename.lower().endswith((".wav", ".mp3", ".flac", ".ogg", ".m4a")):
        raise HTTPException(status_code=400, detail="Unsupported file format. Please upload a WAV, MP3, FLAC, OGG, or M4A file.")

    # Save uploaded file
    file_id = uuid.uuid4().hex[:8]
    save_path = UPLOAD_DIR / f"{file_id}_{file.filename}"
    with open(save_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    try:
        # Step 1: Normalize audio
        logger.info(f"Normalizing audio: {save_path}")
        norm_path = normalize_audio(str(save_path))
        logger.info(f"Normalised audio at: {norm_path}")

        result_map = {"audio_input": file.filename}

        # Step 2: Transcribe
        logger.info("Transcribing audio...")
        text = transcribe_audio(norm_path)
        result_map["audio_text"] = text
        logger.info(f"Transcription: {text}")

        # Step 3: Text emotion (RoBERTa)
        text_emotion, text_prob = predict_text_emotion(text)
        result_map["roberta_text_emotion"] = text_emotion
        result_map["roberta_text_probability"] = round(text_prob, 4) if text_emotion else None

        # Step 4: Audio emotion models
        result_map["wav2vec2_audio_emotion"] = predict_wav2vec2_emotion(norm_path)
        result_map["cnn_audio_emotion"] = predict_cnn_emotion(norm_path)
        result_map["svm_audio_emotion"] = predict_svm_emotion(norm_path)
        result_map["rf_audio_emotion"] = predict_rf_emotion(norm_path)

        # Step 5: Anomaly detection
        report = evaluate_mismatch(result_map)

        return JSONResponse(content=report)

    except Exception as e:
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
    finally:
        # Cleanup uploaded files
        if save_path.exists():
            save_path.unlink()


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False)