app.py

import os
import tempfile
import whisperx
from pyannote.audio import Pipeline
import pandas as pd
import librosa
import soundfile as sf
import numpy as np
from scipy.signal import butter, filtfilt
from typing import Optional, Dict, List, Any
import torch
from dataclasses import dataclass, field
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import time
import shutil
from starlette.concurrency import run_in_threadpool
import gc
from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor
try:
    import noisereduce as nr
    HAVE_NOISEREDUCE = True
except ImportError:
    HAVE_NOISEREDUCE = False
Annotation: Any = None
Segment: Any = None

device = "cuda" if torch.cuda.is_available() else "cpu"
COMPUTE_TYPE = "float16" if device == "cuda" else "float32"
BATCH_SIZE = 4
token = os.environ.get("HF_TOKEN")
global_diarizer = None
def load_pyannote_pipeline():
    """Loads and returns the Pyannote Diarization pipeline."""
    if not token:
        print("HF_TOKEN not set. Diarization is unavailable.")
        return None
    try:
        pyannote_device = torch.device(device)
        pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1", 
            use_auth_token=token
        ).to(pyannote_device)
        print("Pyannote pipeline loaded dynamically.")
        return pipeline
    except Exception as e:
        print(f"Error loading pyannote pipeline dynamically: {type(e).__name__}: {e}. Diarization will be skipped.")
        return None
model_name = "large-v2"
ALIGN_MODEL_MAP = {
    "ur": "kingabzpro/wav2vec2-large-xls-r-300m-Urdu"}
global_align_model_cache = {}
processor = AutoFeatureExtractor.from_pretrained("facebook/mms-lid-4017")
model = Wav2Vec2ForSequenceClassification.from_pretrained("facebook/mms-lid-4017")
model.to("cpu")
class TimelineItem(BaseModel):
    start: float
    end: float
    speaker: str | None = None
    text: str
class AnalysisResult(BaseModel):
    duration: float
    language: str
    der: float | None = None
    speaker_error: float | None = None
    missed_speech: float | None = None
    false_alarm: float | None = None
    timeline_data: list[TimelineItem]
    raw_transcription: str
    warnings: list[str] = field(default_factory=list)

app = FastAPI(title="Audio Analyzer Backend")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://frontend-audio-analyzer.vercel.app"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
@dataclass
class AnalysisResults:
    timelineData: List[Dict[str, Any]] = field(default_factory=list)
    duration: float = 0.0
    languageCode: str = "unknown"
    diarizationErrorRate: Optional[float] = None
    speakerError: Optional[float] = None
    missedSpeech: Optional[float] = None
    falseAlarm: Optional[float] = None
    warnings: List[str] = field(default_factory=list)
    success: bool = False
    message: str = "Analysis initiated."
    rawTranscriptionText: str = ""
def warn(results: AnalysisResults, code: str, detail: str) -> None:
    msg = f"{code}: {detail}"
    if msg not in results.warnings:
        results.warnings.append(msg)
def set_message(results: AnalysisResults, msg: str) -> None:
    initial_message = "Analysis initiated."
    if results.message and results.message != initial_message:
        results.message += f" | {msg}"
    else:
        results.message = msg
def normalize_speaker(lbl: str) -> str:
    lbl_str = str(lbl)
    return lbl_str.replace("SPEAKER_", "Speaker_").replace("speaker_", "Speaker_")
def temp_wav_path() -> str:
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        return f.name
def force_float(value: Optional[Any]) -> Optional[float]:
    """Ensures value is a native Python float or None. Returns None for NaN/Inf."""
    if value is None:
        return None
    try:
        f_val = float(value)
        if np.isnan(f_val) or np.isinf(f_val):
            return None
        return f_val
    except (TypeError, ValueError, AttributeError):
        return None

def butter_filter(y, sr, lowpass=None, highpass=None, order=4):
    nyq = 0.5 * sr
    if highpass and highpass > 0 and highpass < nyq:
        b, a = butter(order, highpass / nyq, btype="highpass", analog=False)
        y = filtfilt(b, a, y)
    if lowpass and lowpass > 0 and lowpass < nyq:
        b, a = butter(order, lowpass / nyq, btype="lowpass", analog=False)
        y = filtfilt(b, a, y)
    return y

def rms_normalize(y, target_rms=0.8, eps=1e-6):
    rms = (y**2).mean() ** 0.5
    if rms < eps:
        return y
    gain = target_rms / (rms + eps)
    return y * gain

def preprocess_audio(input_path,
                     target_sr=16000,
                     normalize_rms=True,
                     target_rms=0.08,
                     denoise=False,
                     highpass=None,
                     lowpass=None,
                     output_subtype="PCM_16",
                     verbose=False) -> str:
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input audio not found: {input_path}")
    output_path = temp_wav_path()
    y_stereo, sr = sf.read(input_path, dtype='float32')
    if y_stereo.ndim > 1:
        y = librosa.to_mono(y_stereo.T)
    else:
        y = y_stereo
    if sr != target_sr:
        y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
        sr = target_sr
    if highpass or lowpass:
        y = butter_filter(y, sr, highpass=highpass, lowpass=lowpass)
    if denoise and HAVE_NOISEREDUCE:
        try:
            noise_len = int(min(len(y), int(0.5 * sr)))
            noise_clip = y[:noise_len]
            y = nr.reduce_noise(y=y, sr=sr, y_noise=noise_clip, prop_decrease=0.9, verbose=False)
        except Exception:
            pass
    if normalize_rms:
        y = rms_normalize(y, target_rms=target_rms)
    sf.write(output_path, y, sr, subtype=output_subtype)
    return output_path
def analyze_audio(audio_file: str, 
                  preprocess: bool = True, 
                  preprocess_params: Optional[Dict[str, Any]] = None) -> AnalysisResults:
    global global_align_model_cache, ALIGN_MODEL_MAP
    global COMPUTE_TYPE
    global BATCH_SIZE
    results = AnalysisResults()
    ends: List[float] = []      
    rows: List[Dict[str, Any]] = []
    rawTranscriptionText: str = ""
    if device == "cpu":
        num_cores = os.cpu_count() or 4
        print(f"Setting PyTorch threads to {num_cores} for CPU performance optimization.")
        try:
            torch.set_num_threads(num_cores)
            torch.jit.enable_onednn_fusion(True) 
        except Exception as e:
            print(f"Warning: Failed to set PyTorch performance flags: {e}")
    if not os.path.exists(audio_file):
        results.message = f"Error: Input audio file '{audio_file}' not found."
        return results
    audio_for_model = audio_file
    temp_preproc = None
    if preprocess:
        params = {
            "target_sr": 16000, "normalize_rms": True, "target_rms": 0.08,
            "denoise": False, "highpass": None, "lowpass": None,
            "output_subtype": "PCM_16", "verbose": False
        }
        if isinstance(preprocess_params, dict):
            params.update(preprocess_params)
        if params.get("denoise") and not HAVE_NOISEREDUCE:
            warn(results, "DENOISE_SKIP", "Denoise requested but noisereduce not installed; skipping denoise.")
            params["denoise"] = False
        try:
            temp_preproc = preprocess_audio(audio_file, **params)
            audio_for_model = temp_preproc
        except Exception as e:
            warn(results, "PREP_FAIL", f"Preprocessing failed: {e}. Falling back to original audio.")
            audio_for_model = audio_file
            temp_preproc = None
            
    start_ml_time = time.time()
    model = None
    audio_loaded = None
    diarization_pipeline = None
    try:
        print(f"Loading Whisper model '{model_name}' on {device}...")
        model = whisperx.load_model(model_name, device, compute_type="float32")
        audio_loaded = whisperx.load_audio(audio_for_model) 
        print("Detecting language...")
        inputs = processor(audio_loaded, sampling_rate=target_sr, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs).logits
        lang_id = torch.argmax(outputs, dim=-1)[0].item()
        detected_language = model.config.id2label[lang_id]
        languageCode = detected_language
        # lang_result = model.transcribe(audio_loaded, batch_size=4, language=None)
        # language_code_detected = lang_result.get("language") or lang_result.get("detected_language")
        # languageCode = language_code_detected
        # results.languageCode = languageCode
        print("Transcribing audio...")
        transcribed_language = "ur"
        result = model.transcribe(audio_loaded, batch_size=BATCH_SIZE, language= transcribed_language 
          )
        full_text = " ".join([seg['text'] for seg in result.get("segments", [])]).strip()
        results.rawTranscriptionText = full_text
        aligned = {"segments": result["segments"]} 
        print(f"Detected language: {languageCode}. Aligning transcription...")
        aligner_lookup_language = transcribed_language
        align_model = None
        metadata = None
        if aligner_lookup_language not in global_align_model_cache:
            align_model_name = ALIGN_MODEL_MAP.get(aligner_lookup_language)
            try:
                align_model, metadata = whisperx.load_align_model(
                                                  language_code=aligner_lookup_language, 
                                                  model_name=align_model_name, 
                                                  device=device
                )
                global_align_model_cache[aligner_lookup_language] = (align_model, metadata)
                print(f"Alignment model successfully loaded/cached for language: {aligner_lookup_language}")
            except Exception as e:
                warn(results, "ALIGN_LOAD_FAIL", f"Failed to load alignment model for {aligner_lookup_language}: {type(e).__name__}: {e}. Alignment skipped.")
                global_align_model_cache[aligner_lookup_language] = (None, None)
        else:
            align_model, metadata = global_align_model_cache[aligner_lookup_language]
            if align_model:
                print(f"Alignment model loaded from cache for language: {aligner_lookup_language}")
        if align_model:
            try:
                print("Performing word-level alignment...")
                aligned = whisperx.align(
                    result["segments"], 
                    align_model, 
                    metadata, 
                    audio_loaded, 
                    device
                )
            except Exception as e:
                warn(results, "ALIGN_RUN_FAIL", f"Alignment execution failed: {type(e).__name__}: {e}. Using raw segments.")
        else:
            warn(results, "ALIGN_SKIP", "Alignment model unavailable; using raw Whisper segments.")
        print("Cleaning up Whisper model memory...")
        del model
        model = None
        del audio_loaded
        audio_loaded = None
        if device == "cuda":
            torch.cuda.empty_cache() 
        gc.collect() 
        print("Memory cleanup complete.")
        diarize_output = None
        diarization_pipeline = load_pyannote_pipeline()
        if diarization_pipeline is not None:
            print("Performing speaker diarization (Requires HF_TOKEN)...")
            try:
                diarize_output = diarization_pipeline(audio_for_model)
                for segment, _, label in diarize_output.itertracks(yield_label=True):
                    print(f"start={segment.start:.1f}s stop={segment.end:.1f}s {label}")
            except Exception as e:
                warn(results, "DIAR_SKIP", f"Error during diarization (likely token/model failure): {type(e).__name__}: {e}. Skipping diarization.")
                diarize_output = None
        else:
            warn(results, "DIAR_SKIP", "HF_TOKEN not set or Diarization Pipeline failed to load globally. Skipping speaker diarization.")
        if diarization_pipeline is not None:
             print("Cleaning up Pyannote model memory...")
             del diarization_pipeline
             diarization_pipeline = None
             if device == "cuda":
                torch.cuda.empty_cache() 
             gc.collect()
             print("Pyannote cleanup complete.")
        print("Assigning speakers to words...")   
        try:
            diarize_segments_for_assignment = []
            if diarize_output is not None and hasattr(diarize_output, "itertracks"):
                for segment, _, label in diarize_output.itertracks(yield_label=True):
                    diarize_segments_for_assignment.append({
                        "start": float(segment.start),
                        "end": float(segment.end),
                        "speaker": normalize_speaker(label)
                    })
                print(f"DEBUG: Converted {len(diarize_segments_for_assignment)} diarization segments.")
            if diarize_segments_for_assignment:
                diarize_df = pd.DataFrame(diarize_segments_for_assignment)
                final = whisperx.assign_word_speakers(diarize_df, aligned)
            else:
                warn(results, "ASSIGN_FAIL", "Diarization segments were empty or unavailable. Defaulting all to Speaker_1.")
                final = aligned
                for seg in final.get("segments", []):
                    seg["speaker"] = "Speaker_1"
        except Exception as e:
            warn(results, "ASSIGN_SPEAKERS_ERROR", f"Error assigning speakers: {type(e).__name__}: {e}. Falling back to unassigned segments.")
            final = aligned
            for seg in final.get("segments", []):
                seg["speaker"] = "Speaker_1"
        def _get_time_field(d: Dict[str, Any], keys: List[str]) -> Optional[float]:
            """Try multiple possible keys and coerce to native float, returning None if not possible."""
            for k in keys:
                if k in d:
                    try:
                        v = d[k]
                        if v is None:
                            continue
                        f = float(v)
                        if np.isnan(f) or np.isinf(f):
                            return None
                        return f
                    except (TypeError, ValueError):
                        continue
            return None
        for seg in final.get("segments", []):
            seg_speaker = normalize_speaker(seg.get("speaker") or seg.get("speaker_label") or "Speaker_1")
            word_list = seg.get("words") or seg.get("tokens") or seg.get("items") or []
            if not word_list:
                word_start = _get_time_field(seg, ["start", "s", "timestamp", "t0"])
                word_end = _get_time_field(seg, ["end", "e", "t1"])
                if word_start is None:
                    continue
                if word_end is None:
                    word_end = word_start
                rows.append({
                    "start": float(word_start),
                    "end": float(word_end),
                    "text": str(seg.get("text", "")).strip(),
                    "speaker": str(seg_speaker),
                })
                continue
            for w in word_list:
                if not isinstance(w, dict):
                    continue
                word_start = _get_time_field(w, ["start", "s", "timestamp", "t0"])
                word_end = _get_time_field(w, ["end", "e", "t1"])
                if word_start is None:
                    word_start = _get_time_field(seg, ["start", "s"])
                if word_end is None:
                    word_end = _get_time_field(seg, ["end", "e"])          
                if word_start is None:
                    continue
                if word_end is None:
                    word_end = word_start
                word_speaker = normalize_speaker(w.get("speaker") or seg_speaker)
                word_text = (w.get("text") or w.get("word") or w.get("label") or "").strip()
                rows.append({
                    "start": float(word_start),
                    "end": float(word_end),
                    "text": str(word_text),
                    "speaker": str(word_speaker),
                })
        rows = sorted(rows, key=lambda r: r.get("start", 0.0))
        results.timelineData = rows
        for w in rows:
            e = w.get("end")
            f_e = force_float(e)
            if f_e is not None:
                ends.append(f_e)  
    except Exception as e:
        results.message = f"Error during ML processing: {type(e).__name__}: {e}"
        return results
    finally:
        if temp_preproc and os.path.exists(temp_preproc):
            os.remove(temp_preproc)
        results.duration = force_float(max(ends) if ends else 0.0) or 0.0
        end_ml_time = time.time()
        print(f"ML Processing finished in {end_ml_time - start_ml_time:.2f} seconds.")
    results.success = True
    return results
@app.post("/upload", response_model=AnalysisResult)
async def upload_file(audio_file: UploadFile = File(...)):
    start_time = time.time()
    audio_path: Optional[str] = None
    try:
        print("Incoming upload:", getattr(audio_file, "filename", None))

        suffix = audio_file.filename.split(".")[-1] if audio_file.filename else "tmp"
        with tempfile.NamedTemporaryFile(suffix=f".{suffix}", delete=False) as tmp_audio:
            shutil.copyfileobj(audio_file.file, tmp_audio)
            audio_path = tmp_audio.name
        print(f"Received audio file: {audio_file.filename} (saved to {audio_path}), size: {os.path.getsize(audio_path)} bytes")
        preprocessing_config = {"denoise": False}
        print(f"Starting ML processing with audio: {audio_path}, preprocess_params: {preprocessing_config}")
        
        analysis_result = await run_in_threadpool(
            analyze_audio,
            audio_file=audio_path,
            preprocess_params=preprocessing_config
        )
        print("MESSAGE:", analysis_result.message)
        if not analysis_result.success:
            raise HTTPException(status_code=500, detail=analysis_result.message)
        print("DURATION BEFORE RETURN:", analysis_result.duration)
        if analysis_result.duration is None:
            analysis_result.duration = 0.0
        return AnalysisResult(
            duration=force_float(analysis_result.duration) or 0.0,
            language=analysis_result.languageCode,
            timeline_data=[
                TimelineItem(
                    start=force_float(item.get('start')) or 0.0,
                    end=force_float(item.get('end')) or 0.0,
                    speaker=str(item.get('speaker')) if item.get('speaker') else None,
                    text=str(item.get('text', ""))
                ) for item in analysis_result.timelineData
            ],
            raw_transcription=analysis_result.rawTranscriptionText,
            warnings=analysis_result.warnings
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Unexpected error during upload process: {type(e).__name__}: {e}")
    finally:
        if audio_path and os.path.exists(audio_path):
            os.remove(audio_path)
        end_time = time.time()
        print(f"API Request processed in {end_time - start_time:.2f} seconds.")
@app.get("/")
def root():
    return {"message": "Audio Analyzer Backend is running."}