backend/main.py

"""FastAPI app: upload a WAV, get diarized transcription + per-speaker gender."""

from __future__ import annotations

import os
import time
import uuid
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import librosa
import numpy as np
import soundfile as sf
from dotenv import load_dotenv
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles

from . import gender
from .transcribe import Transcription, transcribe_with_diarization

load_dotenv()

BASE_DIR = Path(__file__).resolve().parent.parent
UPLOAD_DIR = Path(os.getenv("UPLOAD_DIR") or (BASE_DIR / "uploads"))
FRONTEND_DIR = BASE_DIR / "frontend"
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

app = FastAPI(title="Audio Gender + Transcription POC")


@app.on_event("startup")
def _warmup_model() -> None:
    # Load the gender model ahead of the first request (best-effort).
    try:
        gender.warmup()
    except Exception:
        pass


@app.get("/", response_class=HTMLResponse)
def index() -> HTMLResponse:
    return HTMLResponse((FRONTEND_DIR / "index.html").read_text(encoding="utf-8"))


@app.get("/api/audio/{file_id}")
def get_audio(file_id: str) -> FileResponse:
    # Guard against path traversal; we only ever serve from UPLOAD_DIR.
    safe_id = Path(file_id).name
    path = UPLOAD_DIR / safe_id
    if not path.exists():
        raise HTTPException(status_code=404, detail="Audio not found")
    return FileResponse(path, media_type="audio/wav")


@app.post("/api/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    api_key: Optional[str] = Form(default=None),
    num_speakers: int = Form(default=int(os.getenv("NUM_SPEAKERS", "2"))),
):
    key = (api_key or os.getenv("SARVAM_API_KEY") or "").strip()
    if not key:
        raise HTTPException(
            status_code=400,
            detail="Server is missing the SARVAM_API_KEY. Set it in the .env file and restart.",
        )

    file_id = f"{uuid.uuid4().hex}.wav"
    dest = UPLOAD_DIR / file_id
    dest.write_bytes(await file.read())

    started_at = datetime.now(timezone.utc)
    t0 = time.perf_counter()

    try:
        result = transcribe_with_diarization(str(dest), key, num_speakers=num_speakers)
    except Exception as exc:  # surface a clean error to the UI
        raise HTTPException(status_code=502, detail=str(exc))

    t1 = time.perf_counter()
    speakers = _detect_genders(str(dest), result)
    t2 = time.perf_counter()
    finished_at = datetime.now(timezone.utc)

    timing = {
        "started_at": started_at.isoformat(),
        "finished_at": finished_at.isoformat(),
        "transcription_seconds": round(t1 - t0, 2),
        "gender_seconds": round(t2 - t1, 2),
        "total_seconds": round(t2 - t0, 2),
    }

    return {
        "file_id": file_id,
        "audio_url": f"/api/audio/{file_id}",
        "language_code": result.language_code,
        "full_transcript": result.full_transcript,
        "timing": timing,
        "speakers": speakers,
        "segments": [
            {
                "start": round(s.start, 2),
                "end": round(s.end, 2),
                "speaker_id": s.speaker_id,
                "text": s.text,
                "gender": speakers.get(s.speaker_id, {}).get("label", "uncertain"),
            }
            for s in result.segments
        ],
    }


def _detect_genders(audio_path: str, result: Transcription) -> dict:
    """Estimate gender per speaker by pooling all of their voiced audio."""
    samples, sr = librosa.load(audio_path, sr=16000, mono=True)

    by_speaker: dict[str, list[np.ndarray]] = defaultdict(list)
    for seg in result.segments:
        if seg.end <= seg.start:
            continue
        start_idx = max(0, int(seg.start * sr))
        end_idx = min(len(samples), int(seg.end * sr))
        if end_idx > start_idx:
            by_speaker[seg.speaker_id].append(samples[start_idx:end_idx])

    speakers: dict[str, dict] = {}
    for speaker_id, chunks in by_speaker.items():
        pooled = np.concatenate(chunks) if chunks else np.array([])
        res = gender.classify(pooled, sr)
        speakers[speaker_id] = {
            "label": res.label,
            "confidence": res.confidence,
            "probs": res.probs,
            "audio_seconds": res.audio_seconds,
        }

    # Ensure every diarized speaker has an entry even with no usable audio.
    for seg in result.segments:
        speakers.setdefault(
            seg.speaker_id,
            {"label": "uncertain", "confidence": 0.0, "probs": {}, "audio_seconds": 0.0},
        )
    return speakers


if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")