Update app.py

app.py

@@ -1,20 +1,19 @@
 """
 Boolook - 음성 기반 감정 분석 책 추천 (HuggingFace Spaces)
 수정사항:
-  - 임베딩 로딩을 백그라운드 스레드로 분리 (타임아웃 방지)
-  - 배치 크기 128로 증가 (속도 향상)
-  - 서버가 먼저 열린 뒤 데이터 로딩 진행
-  - 추천 결과 출력을 JSON 형식으로 단순화
-  - emotion_score: 주감정 단일 수치
-  - user_input / recommendation_books 키 사용
-  - 오디오 type="filepath" + soundfile 분기 처리
-  - 피드백 UI 제거 → /api/feedback 엔드포인트로 대체
+  - final_emotion_model_v3.pth (ResNet-SE + BiLSTM + Attention) 커스텀 모델 통합
+  - superb/wav2vec2-base-superb-er 대신 커스텀 모델로 음성 감정 분류
+  - 모델 클래스 정의 (SEBlock, ResBlock, AttentionPooling, EmotionResNet) 포함
+  - Mel-spectrogram 전처리 + TTA(n_tta=8) 추론 + temperature scaling 적용
+  - 4클래스(Angry/Happy/Neutral/Sad) → 한국어 감정 레이블 매핑
 """
 
 import gradio as gr
 import pandas as pd
 import numpy as np
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import pickle
 import csv
 import json
@@ -35,13 +34,21 @@ logger = logging.getLogger(__name__)
 # ============================================================
 # 설정
 # ============================================================
-BOOK_DB_PATH     = Path("book_db_final.csv")
-FEEDBACK_PATH    = Path("user_feedback.csv")
-SBERT_CACHE_PATH = Path("book_embeddings.pkl")
-SAMPLE_RATE      = 16000
+BOOK_DB_PATH        = Path("book_db_final.csv")
+FEEDBACK_PATH       = Path("user_feedback.csv")
+SBERT_CACHE_PATH    = Path("book_embeddings.pkl")
+EMOTION_MODEL_PATH  = Path("final_emotion_model_v3.pth")
+SAMPLE_RATE         = 16000
 MAX_EMBEDDING_BATCH = 128
 
+# Mel-spectrogram 파라미터 (학습 시 사용한 값과 동일하게 맞출 것)
+N_MELS    = 64
+N_FFT     = 1024
+HOP_LEN   = 512
+MAX_FRAMES = 128   # 시간 축 고정 길이
+
 device = 0 if torch.cuda.is_available() else -1
+torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"디바이스: {'GPU' if device == 0 else 'CPU'}")
 
 # ============================================================
@@ -53,7 +60,245 @@ _data_ready     = False
 _data_lock      = threading.Lock()
 
 # ============================================================
-# 모델 로딩
+# ① 커스텀 감정 모델 아키텍처 정의
+# ============================================================
+
+class SEBlock(nn.Module):
+    """Squeeze-and-Excitation Block"""
+    def __init__(self, channels: int, reduction: int = 16):
+        super().__init__()
+        self.excitation = nn.Sequential(
+            nn.Linear(channels, channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channels // reduction, channels, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        # x: (B, C, H, W)
+        b, c, _, _ = x.shape
+        w = x.mean(dim=[2, 3])          # global avg pool
+        w = self.excitation(w).view(b, c, 1, 1)
+        return x * w
+
+
+class ResBlock(nn.Module):
+    """ResNet Basic Block with SE"""
+    def __init__(self, in_ch: int, out_ch: int, stride: int = 1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False)
+        self.bn1   = nn.BatchNorm2d(out_ch)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
+        self.bn2   = nn.BatchNorm2d(out_ch)
+        self.se    = SEBlock(out_ch, reduction=max(1, out_ch // 16))
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_ch != out_ch:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_ch),
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        out = self.bn2(self.conv2(out))
+        out = self.se(out)
+        out = F.relu(out + self.shortcut(x), inplace=True)
+        return out
+
+
+class AttentionPooling(nn.Module):
+    """Temporal Attention Pooling"""
+    def __init__(self, hidden: int):
+        super().__init__()
+        self.attn = nn.Linear(hidden, 1)
+
+    def forward(self, x):
+        # x: (B, T, H)
+        w = torch.softmax(self.attn(x), dim=1)   # (B, T, 1)
+        return (x * w).sum(dim=1)                 # (B, H)
+
+
+class EmotionResNet(nn.Module):
+    """
+    ResNet-SE + 2-layer BiLSTM + Attention Pooling + Classifier
+    입력: (B, 1, N_MELS, T) Mel-spectrogram
+    출력: (B, num_classes) logits
+    """
+    def __init__(self, num_classes: int = 4):
+        super().__init__()
+        # CNN stem
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(1, 64, 7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+        )
+        # ResNet layers
+        self.layer1 = nn.Sequential(ResBlock(64,  64),  ResBlock(64,  64))
+        self.layer2 = nn.Sequential(ResBlock(64,  128, stride=2), ResBlock(128, 128))
+        self.layer3 = nn.Sequential(ResBlock(128, 256, stride=2), ResBlock(256, 256))
+
+        # BiLSTM (2 layers)
+        self.bilstm = nn.LSTM(
+            input_size=256, hidden_size=256,
+            num_layers=2, batch_first=True,
+            bidirectional=True, dropout=0.3,
+        )
+
+        # Attention
+        self.attention = AttentionPooling(hidden=512)
+
+        # Classifier
+        self.classifier = nn.Sequential(
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.5),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, x):
+        # CNN
+        x = F.relu(self.conv1(x), inplace=True)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        # (B, C, H, W) → temporal sequence: global-avg over freq axis
+        x = x.mean(dim=2)          # (B, C, W)
+        x = x.permute(0, 2, 1)    # (B, T, C)
+
+        # BiLSTM
+        x, _ = self.bilstm(x)     # (B, T, 512)
+
+        # Attention pooling
+        x = self.attention(x)     # (B, 512)
+
+        return self.classifier(x)
+
+
+# ============================================================
+# ② 커스텀 감정 모델 로드
+# ============================================================
+_emotion_model      = None
+_emotion_classes    = ["Angry", "Happy", "Neutral", "Sad"]
+_emotion_label_enc  = None
+_emotion_temp       = 1.0
+_emotion_n_tta      = 1
+
+def _load_emotion_model():
+    global _emotion_model, _emotion_classes, _emotion_label_enc, _emotion_temp, _emotion_n_tta
+    if not EMOTION_MODEL_PATH.exists():
+        logger.error(f"{EMOTION_MODEL_PATH} 파일이 없습니다. 커스텀 감정 모델을 사용하지 않습니다.")
+        return
+
+    try:
+        ckpt = torch.load(EMOTION_MODEL_PATH, map_location="cpu", weights_only=False)
+
+        _emotion_classes   = [str(c) for c in ckpt.get("classes", _emotion_classes)]
+        _emotion_label_enc = ckpt.get("label_encoder", None)
+        _emotion_temp      = float(ckpt.get("temperature", 1.0))
+        _emotion_n_tta     = int(ckpt.get("n_tta", 1))
+
+        model = EmotionResNet(num_classes=len(_emotion_classes))
+        model.load_state_dict(ckpt["model_state_dict"])
+        model.to(torch_device)
+        model.eval()
+
+        _emotion_model = model
+        logger.info(
+            f"커스텀 감정 모델 로드 완료 | "
+            f"클래스: {_emotion_classes} | "
+            f"val_acc: {ckpt.get('val_accuracy', 'N/A')} | "
+            f"val_f1: {ckpt.get('best_val_f1', 'N/A'):.4f} | "
+            f"temp: {_emotion_temp} | TTA: {_emotion_n_tta}"
+        )
+    except Exception as e:
+        logger.error(f"커스텀 감정 모델 로드 실패: {e}")
+
+_load_emotion_model()
+
+# ============================================================
+# ③ Mel-spectrogram 전처리
+# ============================================================
+def _compute_melspec(y: np.ndarray, sr: int) -> torch.Tensor:
+    """
+    오디오 배열 → (1, 1, N_MELS, MAX_FRAMES) 텐서
+    librosa 없이 torch만 사용하는 간이 구현
+    """
+    try:
+        import librosa
+        mel = librosa.feature.melspectrogram(
+            y=y, sr=sr,
+            n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP_LEN,
+        )
+        mel_db = librosa.power_to_db(mel, ref=np.max)
+    except ImportError:
+        # librosa 없을 때 torch STFT로 대체
+        wav = torch.tensor(y, dtype=torch.float32)
+        window = torch.hann_window(N_FFT)
+        stft = torch.stft(wav, N_FFT, HOP_LEN, return_complex=True, window=window)
+        power = stft.abs() ** 2                         # (freq, T)
+        # 간이 mel filterbank (삼각형 근사)
+        mel_fb = torch.zeros(N_MELS, power.shape[0])
+        for m in range(N_MELS):
+            mel_fb[m, m * (power.shape[0] // N_MELS):
+                       (m + 1) * (power.shape[0] // N_MELS)] = 1.0
+        mel = mel_fb @ power                            # (N_MELS, T)
+        mel_db = (mel + 1e-6).log().numpy()
+
+    # 정규화
+    mel_db = (mel_db - mel_db.mean()) / (mel_db.std() + 1e-6)
+
+    # 시간 축 패딩/자르기
+    T = mel_db.shape[1]
+    if T < MAX_FRAMES:
+        mel_db = np.pad(mel_db, ((0, 0), (0, MAX_FRAMES - T)), mode="constant")
+    else:
+        mel_db = mel_db[:, :MAX_FRAMES]
+
+    # (1, 1, N_MELS, MAX_FRAMES)
+    tensor = torch.tensor(mel_db, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+    return tensor.to(torch_device)
+
+
+# ============================================================
+# ④ TTA 추론
+# ============================================================
+def _tta_augment(spec: torch.Tensor) -> torch.Tensor:
+    """단순 시간 이동 augmentation for TTA"""
+    shift = np.random.randint(-MAX_FRAMES // 8, MAX_FRAMES // 8)
+    return torch.roll(spec, shift, dims=-1)
+
+
+def _infer_emotion_model(y: np.ndarray, sr: int) -> Dict[str, float]:
+    """커스텀 모델 추론 → 클래스별 확률 dict (원본 영�� 레이블)"""
+    if _emotion_model is None:
+        return {c: 0.0 for c in _emotion_classes}
+
+    try:
+        spec = _compute_melspec(y, sr)   # (1, 1, N_MELS, T)
+
+        logits_list = []
+        with torch.no_grad():
+            n = max(1, _emotion_n_tta)
+            for i in range(n):
+                inp = _tta_augment(spec) if i > 0 else spec
+                logits = _emotion_model(inp)              # (1, num_classes)
+                logits_list.append(logits)
+
+        avg_logits = torch.stack(logits_list).mean(dim=0)          # (1, C)
+        probs = torch.softmax(avg_logits / _emotion_temp, dim=-1)  # temperature scaling
+        probs = probs[0].cpu().numpy()
+
+        return {cls: float(p) for cls, p in zip(_emotion_classes, probs)}
+
+    except Exception as e:
+        logger.error(f"커스텀 모델 추론 실패: {e}")
+        return {c: 0.0 for c in _emotion_classes}
+
+
+# ============================================================
+# 모델 로딩 (STT, SBERT)
 # ============================================================
 logger.info("모델 로딩 중...")
 
@@ -79,17 +324,6 @@ try:
 except Exception as e:
     logger.error(f"SBERT 로드 실패: {e}")
 
-audio_emotion_pipeline = None
-try:
-    audio_emotion_pipeline = hf_pipeline(
-        "audio-classification",
-        model="superb/wav2vec2-base-superb-er",
-        device=device,
-    )
-    logger.info("음성 감정 모델 로드 완료")
-except Exception as e:
-    logger.warning(f"음성 감정 모델 스킵 (텍스트만 사용): {e}")
-
 logger.info("모델 로딩 완료!")
 
 # ============================================================
@@ -118,7 +352,13 @@ if sbert_model:
     except Exception as e:
         logger.error(f"감정 레이블 임베딩 실패: {e}")
 
-_AUDIO_LABEL_MAP = {"hap": "기쁨", "neu": "신뢰", "sad": "슬픔", "ang": "분노"}
+# 커스텀 모델 영어 레이블 → 한국어 매핑
+_CUSTOM_LABEL_MAP: Dict[str, str] = {
+    "Happy":   "기쁨",
+    "Sad":     "슬픔",
+    "Angry":   "분노",
+    "Neutral": "신뢰",
+}
 
 _KEYWORD_BOOSTS = {
     "슬픔": ["슬프", "우울", "눈물", "힘들", "외로"],
@@ -246,25 +486,23 @@ def text_emotion_scores(text: str) -> Dict[str, float]:
     return scores
 
 
-def audio_emotion_scores(audio_array: np.ndarray, sr: int) -> Dict[str, float]:
-    scores = {emo: 0.0 for emo in _EMOTION_LABELS}
-    if audio_emotion_pipeline is None:
-        return scores
+def audio_emotion_scores(y: np.ndarray, sr: int) -> Dict[str, float]:
+    """
+    커스텀 모델(final_emotion_model_v3.pth)로 음성 감정 점수 반환.
+    영어 4클래스 확률을 한국어 8클래스 공간으로 매핑.
+    """
+    base = {emo: 0.0 for emo in _EMOTION_LABELS}
 
-    try:
-        import scipy.io.wavfile as wav_io
-        import tempfile
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            wav_io.write(tmp.name, sr, (audio_array * 32767).astype(np.int16))
-            results = audio_emotion_pipeline(tmp.name)
-            Path(tmp.name).unlink(missing_ok=True)
-        for item in results:
-            mapped = _AUDIO_LABEL_MAP.get(item["label"])
-            if mapped:
-                scores[mapped] += item["score"]
-    except Exception as e:
-        logger.warning(f"음성 감정 실패: {e}")
-    return scores
+    raw = _infer_emotion_model(y, sr)   # {"Happy": 0.6, "Sad": 0.2, ...}
+    if not raw or all(v == 0 for v in raw.values()):
+        return base
+
+    for eng_label, prob in raw.items():
+        kor_label = _CUSTOM_LABEL_MAP.get(eng_label)
+        if kor_label and kor_label in base:
+            base[kor_label] += prob
+
+    return base
 
 
 def fused_emotion(t_scores: Dict[str, float], a_scores: Dict[str, float]) -> Tuple[str, Dict[str, float]]:
@@ -280,6 +518,7 @@ def fused_emotion(t_scores: Dict[str, float], a_scores: Dict[str, float]) -> Tup
     top_emotion = max(combined, key=combined.get)
     return top_emotion, combined
 
+
 # ============================================================
 # 추천
 # ============================================================
@@ -324,6 +563,7 @@ def get_recommendations(user_input: str, emotion: str, top_n: int = 3) -> List[D
         logger.error(f"추천 실패: {e}")
         return []
 
+
 # ============================================================
 # 추천 결과 → JSON 렌더링
 # ============================================================
@@ -349,6 +589,7 @@ def _render_books_json(user_input: str, emotion: str, combined: Dict[str, float]
     }
     return json.dumps(output, ensure_ascii=False, indent=2)
 
+
 # ============================================================
 # 피드백
 # ============================================================
@@ -423,12 +664,12 @@ def api_feedback(feedback_data) -> str:
 
 def get_feedback_stats() -> str:
     if not FEEDBACK_PATH.exists():
-        return "📊 아직 피드백이 없습니다."
+        return "아직 피드백이 없습니다."
     try:
         fb_df = pd.read_csv(FEEDBACK_PATH, encoding="utf-8-sig", on_bad_lines="skip")
         total = len(fb_df)
         if total == 0:
-            return "📊 아직 피드백이 없습니다."
+            return "아직 피드백이 없습니다."
         emo_counts = fb_df.groupby("emotion")["accepted"].agg(["count", "sum"])
         lines = [f"**총 피드백: {total}건**\n"]
         for emo, row_s in emo_counts.iterrows():
@@ -440,6 +681,7 @@ def get_feedback_stats() -> str:
     except Exception as e:
         return f"통계 로드 실패: {e}"
 
+
 # ============================================================
 # 메인 처리
 # ============================================================
@@ -477,7 +719,7 @@ def process_voice(audio_input):
             return json.dumps({"error": "음성이 인식되지 않았습니다."}, ensure_ascii=False, indent=2), []
 
         t_scores = text_emotion_scores(user_input)
-        a_scores = audio_emotion_scores(y, sr)
+        a_scores = audio_emotion_scores(y, sr)         # ← 커스텀 모델 사용
         top_label, combined = fused_emotion(t_scores, a_scores)
         books      = get_recommendations(user_input, top_label, top_n=3)
         books_json = _render_books_json(user_input, top_label, combined, books)
@@ -493,40 +735,40 @@ def run_analysis(audio):
     books_json, books = process_voice(audio)
     return books_json, books
 
+
 # ============================================================
 # Gradio UI
 # ============================================================
-with gr.Blocks(theme=gr.themes.Soft(), title="Boolook 📚") as demo:
+with gr.Blocks(theme=gr.themes.Soft(), title="Boolook") as demo:
     gr.Markdown("""
-    # 📚 Boolook — 음성 기반 감정 분석 책 추천
+    # Boolook — 음성 기반 감정 분석 책 추천
     당신의 감정을 말로 표현하면, AI가 딱 맞는 책을 추천해드립니다.
-
-    🎤 **사용법:** 마이크로 감정 표현 → 분석하기
+    **사용법:** 마이크로 감정 표현 → 분석하기
     """)
 
     state_books = gr.State([])
 
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 🎤 음성 입력")
+            gr.Markdown("### 음성 입력")
             audio_in    = gr.Audio(
                 sources=["microphone", "upload"],
                 type="filepath",
                 label="마이크 또는 파일 업로드",
             )
-            analyze_btn = gr.Button("🔍 분석하기", variant="primary", size="lg")
-            gr.Markdown("💡 예: '오늘 너무 슬퍼요', '행복한 기분이에요'")
+            analyze_btn = gr.Button("분석하기", variant="primary", size="lg")
+            gr.Markdown("예: '오늘 너무 슬퍼요', '행복한 기분이에요'")
 
         with gr.Column(scale=1):
             out_books_json = gr.Code(
-                label="📊 분석 결과 & 📖 추천 도서",
+                label="분석 결과 & 추천 도서",
                 language="json",
                 interactive=False,
             )
 
-    with gr.Accordion("📈 통계", open=False):
+    with gr.Accordion("통계", open=False):
         stats_md    = gr.Markdown("새로고침을 눌러주세요.")
-        refresh_btn = gr.Button("🔄 통계 새로고침")
+        refresh_btn = gr.Button("통계 새로고침")
         refresh_btn.click(fn=get_feedback_stats, outputs=stats_md)
 
     # 피드백 API 엔드포인트 (클라이언트 전용, UI 미노출)
@@ -548,4 +790,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Boolook 📚") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()