app.py

import os
import re
import shutil
import tempfile
import traceback
import numpy as np
import librosa
import soundfile as sf
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Depends, Security
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import JSONResponse, PlainTextResponse
from funasr import AutoModel

app = FastAPI(title="OpenAI Compatible STT & Unified Report API")

# ================= 1. 安全认证与配置 =================
API_KEY = os.getenv("API_KEY")
security = HTTPBearer(auto_error=False)

def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
    """校验请求头中的 Bearer Token"""
    if not API_KEY:
        return True
    if not credentials or credentials.credentials != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Unauthorized: Invalid or missing API Key."
        )
    return True

# 亲昵口语热词配置
CUSTOM_HOTWORDS = os.getenv(
    "HOTWORDS", 
    "爹地 妈咪 学长 学姐 欧巴 欧尼 宝贝 亲爱的 哥哥 姐姐 弟弟 妹妹 师傅 老公 老婆 爹咪 小路 小路哥哥 因因 岚岚 甜甜 芝芝瓜 小选 小蛇 宝宝 叶瑄 艾因 路辰 罗夏 司岚 星星嘟嘟 牧首 教宗 虎先生 小叶哥哥"
)

# 口语错别字规则纠错映射
COLLOQUIAL_CORRECTIONS = {
    r"跌地|跌弟|爹底|迭地|嗲地": "爹地",
    r"学掌|鞋长|薛长|雪长": "学长",
    r"学姐|薛姐|雪姐": "学姐",
    r"欧八|偶吧|藕爸|偶霸": "欧巴",
    r"欧妮|偶妮|偶泥": "欧尼",
    r"马咪|妈米|麻咪": "妈咪",
    r"抱抱我|鲍鲍我": "抱抱我"
}

def correct_colloquial_words(text: str) -> str:
    """自动修复口语常见同音错别字"""
    for pattern, target in COLLOQUIAL_CORRECTIONS.items():
        text = re.sub(pattern, target, text)
    return text

def format_cjk_spacing(text: str) -> str:
    """在中文字符与英文字母/数字之间自动插入规范空格"""
    text = re.sub(r'([\u4e00-\u9fa5])([a-zA-Z0-9])', r'\1 \2', text)
    text = re.sub(r'([a-zA-Z0-9])([\u4e00-\u9fa5])', r'\1 \2', text)
    return text

# ================= 2. 加载模型 (优先本地离线) =================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
LOCAL_ASR_PATH = os.path.join(BASE_DIR, "models", "SenseVoiceSmall")
LOCAL_EMO_PATH = os.path.join(BASE_DIR, "models", "emotion2vec_plus_seed")

asr_model_path = LOCAL_ASR_PATH if os.path.exists(LOCAL_ASR_PATH) else "FunAudioLLM/SenseVoiceSmall"
emo_model_path = LOCAL_EMO_PATH if os.path.exists(LOCAL_EMO_PATH) else "emotion2vec/emotion2vec_plus_seed"

print(f"[*] Loading ASR from: {asr_model_path}")
asr_model = AutoModel(
    model=asr_model_path,
    device="cpu",
    disable_update=True
)

print(f"[*] Loading EMO from: {emo_model_path}")
emo_model = AutoModel(
    model=emo_model_path,
    device="cpu",
    disable_update=True
)

# ================= 3. 音频前处理 (无损原声) =================
def preprocess_audio(input_path: str, target_sr: int = 16000):
    """仅做 16kHz 重采样与防爆音峰值归一化，保留完整声学泛音与弱音"""
    y, sr = librosa.load(input_path, sr=target_sr, mono=True)
    if len(y) == 0:
        return input_path, 0.0

    duration = float(len(y) / sr)

    # 峰值归一化至 -0.5 dB
    max_val = float(np.max(np.abs(y)))
    if max_val > 0:
        y = (y / max_val) * 0.95

    processed_path = input_path + "_processed.wav"
    sf.write(processed_path, y, sr, subtype="PCM_16")
    return processed_path, duration

# ================= 4. 声学通用时序停顿与拉长音检测 =================
def detect_prosody_marks(audio_path: str, raw_text: str) -> str:
    """
    通用声学韵律标点检测（不再依赖固定词表，全句型通用）：
    - 任意字词叠音结巴: 我…我 / 这个…这个 / 啊…啊、啊！
    - 短促顿音 (0.15s~0.35s): 、
    - 短语气停顿 (0.35s~0.75s): …
    - 长时思考/留白 (>=0.75s): ……
    - 尾音明显拖长 (>=0.65s): ~
    """
    text = raw_text.strip()
    if not text:
        return text

    # 1. 任意单字/词组叠音卡顿处理 (如 "啊 啊" -> "啊…啊", "这个 这个" -> "这个…这个")
    def replace_stutter(match):
        w1, sep, w2 = match.group(1), match.group(2), match.group(3)
        if w1 == w2:
            return f"{w1}…{w2}"
        return f"{w1}、{w2}"

    text = re.sub(r'([\u4e00-\u9fa5a-zA-Z]{1,2})([,\s，]+)([\u4e00-\u9fa5a-zA-Z]{1,2})', replace_stutter, text)

    # 2. 结合音频物理静音间隙（Silence Gaps）动态映射语气停顿深度
    try:
        y, sr = librosa.load(audio_path, sr=16000, mono=True)
        # 获取有效发音区间 (高于 25dB 视为发音)
        intervals = librosa.effects.split(y, top_db=25)
        
        # 统计句中的最大物理停顿
        if len(intervals) > 1:
            max_gap = 0.0
            for i in range(len(intervals) - 1):
                gap = (intervals[i+1][0] - intervals[i][1]) / sr
                if gap > max_gap:
                    max_gap = gap

            if max_gap >= 0.75:
                # 长停顿 (>= 0.75s) 映射为 ……
                text = re.sub(r'[，,]', '……', text, count=1)
            elif max_gap >= 0.35:
                # 短停顿 (0.35s ~ 0.75s) 映射为 …
                text = re.sub(r'[，,]', '…', text, count=1)
            elif 0.15 <= max_gap < 0.35:
                # 短促顿音 (0.15s ~ 0.35s) 映射为 、
                text = re.sub(r'[，,]', '、', text, count=1)

        # 3. 尾音拉长检测（任意末尾字音节持续发音 >= 0.65s 自动追加 ~）
        if len(intervals) > 0:
            last_start, last_end = intervals[-1]
            last_duration = (last_end - last_start) / sr
            
            if last_duration >= 0.65:
                if not text.endswith("~") and not text.endswith("……") and not text.endswith("…"):
                    text += "~"
    except Exception as e:
        print(f"[Warn] Prosody detection fallback: {e}")

    return text

# ================= 5. 情绪与时序解析 =================
DEFAULT_LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "sad", "surprised", "other"]

def get_energy_state(tone: str) -> str:
    if tone in ["angry", "surprised"]:
        return "high / rising"
    elif tone in ["sad"]:
        return "low / falling"
    elif tone in ["happy", "coquetry"]:
        return "cheerful / dynamic"
    return "moderate / flat"

def parse_emotion_and_timeline(emo_res, total_duration: float, window_size: float = 5.0):
    """提取全局情绪与时序切片（10秒以下不切分）"""
    emotion_scores = {}
    top_emotion = "neutral"
    timeline = []

    if not emo_res or len(emo_res) == 0:
        return "neutral", {"neutral": 1.0}, [{
            "start": 0.0,
            "end": round(float(total_duration), 2),
            "tone": "neutral",
            "confidence": 1.0,
            "energy_state": "moderate / flat"
        }]

    item = emo_res[0]
    labels = item.get("labels", DEFAULT_LABELS)
    scores = item.get("scores", None)

    if isinstance(scores, dict):
        emotion_scores = {str(k): round(float(v), 4) for k, v in scores.items()}
    elif isinstance(scores, (list, np.ndarray)):
        score_list = [float(x) for x in scores]
        for idx, s in enumerate(score_list):
            lbl = labels[idx] if idx < len(labels) else f"label_{idx}"
            emotion_scores[str(lbl)] = round(s, 4)

    if not emotion_scores:
        emotion_scores = {"neutral": 1.0}

    top_emotion = max(emotion_scores, key=emotion_scores.get)

    # 10 秒及以下不切分，直接返回整段单一切片
    if total_duration <= 10.0:
        seg_tone = top_emotion if emotion_scores[top_emotion] >= 0.60 else "neutral"
        timeline.append({
            "start": 0.0,
            "end": round(float(total_duration), 2),
            "tone": seg_tone,
            "confidence": emotion_scores[top_emotion],
            "energy_state": get_energy_state(seg_tone)
        })
        return top_emotion, emotion_scores, timeline

    # 10 秒以上时间轴切片
    raw_feats = item.get("feats") if "feats" in item else None
    if raw_feats is not None:
        try:
            feats = np.array(raw_feats, dtype=np.float32)
            if feats.ndim == 1:
                feats = np.expand_dims(feats, axis=0)
            
            total_frames = feats.shape[0]
            if total_frames > 0:
                frames_per_window = max(1, int(total_frames * (window_size / max(total_duration, 0.1))))
                num_windows = int(np.ceil(total_frames / frames_per_window))

                for i in range(num_windows):
                    start_f = i * frames_per_window
                    end_f = min(total_frames, (i + 1) * frames_per_window)
                    chunk = feats[start_f:end_f]
                    mean_vec = np.mean(chunk, axis=0) if chunk.ndim > 1 else chunk

                    start_t = round(float(i * window_size), 2)
                    end_t = round(float(min(total_duration, (i + 1) * window_size)), 2)

                    if len(mean_vec) == len(labels):
                        t_idx = int(np.argmax(mean_vec))
                        t_label = labels[t_idx]
                        conf = float(mean_vec[t_idx])
                    else:
                        t_label = top_emotion
                        conf = float(emotion_scores.get(top_emotion, 0.8))

                    # 局部时序置信度不足 0.60 同样保底降级为 neutral
                    final_seg_tone = str(t_label) if conf >= 0.60 else "neutral"

                    timeline.append({
                        "start": start_t,
                        "end": end_t,
                        "tone": final_seg_tone,
                        "confidence": round(float(conf), 4),
                        "energy_state": get_energy_state(final_seg_tone)
                    })
        except Exception as e:
            print(f"[Warn] Timeline parse failed: {e}")

    if not timeline:
        timeline.append({
            "start": 0.0,
            "end": round(float(total_duration), 2),
            "tone": top_emotion if emotion_scores[top_emotion] >= 0.60 else "neutral",
            "confidence": emotion_scores[top_emotion],
            "energy_state": get_energy_state(top_emotion if emotion_scores[top_emotion] >= 0.60 else "neutral")
        })

    return top_emotion, emotion_scores, timeline

def evaluate_tone_and_emotion(text: str, emotion_scores: dict, raw_top_emotion: str):
    """
    结合声学与文本语义进行多模态裁决：
    1. 撒娇/软萌判定门槛为 0.45
    2. 常规声学情绪（happy/sad/angry/等）必须 >= 0.60，否则保底为 neutral
    """
    top_score = float(emotion_scores.get(raw_top_emotion, 0.0))

    happy_score = float(emotion_scores.get("happy", 0.0))
    surprised_score = float(emotion_scores.get("surprised", 0.0))
    has_soft_words = bool(re.search(r"[嘛啦呀吧啵呗嘤~～…]", text))

    # 1. 撒娇判定
    if (happy_score + surprised_score >= 0.45) and has_soft_words:
        return "coquetry", max(happy_score, top_score)

    # 2. 常规声学情绪阈值校验
    if top_score < 0.60:
        return "neutral", top_score

    # 3. 满足置信度要求，正式采纳该情绪
    return raw_top_emotion, top_score

# ================= 6. 整合文本报告 =================
EMOTION_ZH_MAP = {
    "neutral": "平静/中性",
    "happy": "开心/愉悦",
    "sad": "低落/悲伤",
    "angry": "愤怒/生气",
    "surprised": "惊讶/上扬",
    "fearful": "害怕/紧张",
    "disgusted": "厌恶/反感",
    "coquetry": "撒娇/软萌"
}

def format_unified_text_report(clean_text: str, final_emotion: str, confidence: float, duration: float, emotion_scores: dict, prosody_timeline: list) -> str:
    """将所有分析维度拼接为一段结构化的纯文本"""
    emo_desc = EMOTION_ZH_MAP.get(final_emotion, final_emotion)
    sorted_scores = sorted(emotion_scores.items(), key=lambda x: x[1], reverse=True)
    scores_str = " | ".join([f"{k}: {round(v * 100, 2)}%" for k, v in sorted_scores if v > 0.01])
    
    timeline_items = []
    for item in prosody_timeline:
        timeline_items.append(f"{item['start']}s-{item['end']}s: {item['tone']} ({item['energy_state']})")
    timeline_str = " -> ".join(timeline_items)

    lines = [
        f"[语音识别文本]：{clean_text}",
        f"[主导语气]：{final_emotion} ({emo_desc})",
        f"[综合置信度]：{round(confidence, 4)}",
        f"[音频时长]：{round(duration, 2)}秒",
        f"[声学情绪分布]：{scores_str}",
        f"[语调起伏走势]：{timeline_str}",
        f"注：识别文本有概率因读音近似出现错误"
    ]
    return "\n".join(lines)

# ================= 7. API 路由端点 =================
@app.get("/")
async def root():
    return {
        "status": "online",
        "service": "SenseVoice + Emotion2Vec+ Unified Report API",
        "docs": "/docs",
        "endpoint": "/v1/audio/transcriptions"
    }

@app.post("/v1/audio/transcriptions")
async def transcribe_audio(
    file: UploadFile = File(...),
    model_name: str = Form(default="whisper-1", alias="model"),
    language: str = Form(default="auto"),
    response_format: str = Form(default="json"),
    authenticated: bool = Depends(verify_token)
):
    suffix = os.path.splitext(file.filename)[-1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        shutil.copyfileobj(file.file, tmp)
        raw_path = tmp.name

    processed_path = None
    try:
        # 1. 音频前处理（无损原声）
        processed_path, duration = preprocess_audio(raw_path)

        # 2. SenseVoice 识别文本
        asr_lang = language if language in ["zh", "en", "yue", "ja", "ko", "auto"] else "auto"
        asr_res = asr_model.generate(
            input=processed_path,
            language=asr_lang,
            use_itn=True,
            hotword=CUSTOM_HOTWORDS
        )
        raw_text = asr_res[0]["text"] if asr_res and len(asr_res) > 0 else ""
        clean_text = re.sub(r"<\|.*?\|>", "", raw_text).strip()
        
        # 3. 声学通用时序停顿与拉长音注入
        clean_text = detect_prosody_marks(processed_path, clean_text)

        # 4. 中文口语纠错与中英空格排版
        clean_text = correct_colloquial_words(clean_text)
        clean_text = format_cjk_spacing(clean_text)

        # 5. Emotion2Vec+ 提取情绪与概率
        emo_res = emo_model.generate(
            input=processed_path,
            granularity="utterance"
        )
        raw_top_emotion, emotion_scores, prosody_timeline = parse_emotion_and_timeline(
            emo_res, 
            duration, 
            window_size=5.0
        )

        # 6. 多模态综合语气裁决（阈值 >= 0.60）
        final_emotion, confidence = evaluate_tone_and_emotion(
            clean_text, 
            emotion_scores, 
            raw_top_emotion
        )

        # 7. 生成整合文本报告
        unified_report_text = format_unified_text_report(
            clean_text=clean_text,
            final_emotion=final_emotion,
            confidence=confidence,
            duration=duration,
            emotion_scores=emotion_scores,
            prosody_timeline=prosody_timeline
        )

        if response_format == "text":
            return PlainTextResponse(unified_report_text)
        elif response_format == "verbose_json":
            return JSONResponse({
                "task": "transcribe",
                "language": str(asr_lang),
                "duration": round(float(duration), 2),
                "text": unified_report_text,
                "raw_text": clean_text,
                "emotion": str(final_emotion),
                "confidence": round(float(confidence), 4),
                "emotion_scores": emotion_scores,
                "prosody_timeline": prosody_timeline,
                "segments": []
            })
        else:
            return JSONResponse({
                "text": unified_report_text,
                "raw_text": clean_text,
                "duration": round(float(duration), 2),
                "emotion": str(final_emotion),
                "confidence": round(float(confidence), 4),
                "emotion_scores": emotion_scores,
                "prosody_timeline": prosody_timeline
            })

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(raw_path):
            os.remove(raw_path)
        if processed_path and os.path.exists(processed_path):
            os.remove(processed_path)