real_time_processor.py

import asyncio
import time
import threading
from queue import Queue
import torch
import numpy as np

from speech_to_text import RealTimeSTT
from audio_segmenter import AudioSegmenter
from models.classifier import EmotionClassifier
from features.extractor_korean import extract_embedding


class RealTimeEmotionProcessor:
    """
    발화 단위 기반 실시간 감정 분석 처리기
    STT 결과를 트리거로 해당 구간의 음성을 분석
    """
    
    def __init__(self):
        # STT 모듈 초기화
        self.stt = RealTimeSTT()
        
        # 음성 세그먼터 초기화
        self.audio_segmenter = AudioSegmenter()
        
        # 감정 분류 모델 로드
        self.emotion_model = EmotionClassifier()
        self.emotion_model.eval()
        # self.emotion_model.load_state_dict(torch.load("outputs/emotion_model.pth"))
        
        # 감정 라벨
        self.emotion_labels = ['행복', '슬픔', '놀람', '분노', '우울']
        
        # 결과 큐 (Late Fusion용)
        self.audio_emotion_queue = Queue()
        
        # 처리 상태
        self.is_processing = False
        
    def start_processing(self):
        """실시간 처리 시작"""
        self.is_processing = True
        
        # STT 콜백 설정 - 발화 완료시 트리거
        self.stt.set_utterance_callback(self.on_utterance_complete)
        
        # 백그라운드에서 음성 녹음 시작
        self.audio_segmenter.start_recording()
        
        # STT 시작
        self.stt.start_listening()
        
        print("🎤 실시간 감정 분석 시작...")
        print("말씀하세요! 발화가 끝나면 감정을 분석합니다.")
        
    def on_utterance_complete(self, text, start_time, end_time):
        """
        발화 완료 시 호출되는 콜백
        해당 시간 구간의 음성을 추출하여 감정 분석
        """
        if not text.strip():
            return
            
        print(f"\n📝 인식된 텍스트: '{text}'")
        print(f"⏱️  시간 구간: {start_time:.2f}s ~ {end_time:.2f}s")
        
        # 해당 구간의 음성 데이터 추출
        audio_segment = self.audio_segmenter.extract_segment(start_time, end_time)
        
        if audio_segment is not None:
            # 음성 감정 분석 수행
            emotion_result = self.analyze_audio_emotion(audio_segment)
            
            # 결과 출력
            self.print_emotion_result(text, emotion_result)
            
            # Late Fusion을 위해 큐에 저장
            self.audio_emotion_queue.put({
                'text': text,
                'start_time': start_time,
                'end_time': end_time,
                'audio_emotion': emotion_result,
                'timestamp': time.time()
            })
        
    def analyze_audio_emotion(self, audio_segment):
        """
        음성 세그먼트에서 감정 분석
        기존 wav2vec2 + MLP 모델 활용
        """
        try:
            # 임시 파일로 저장 (기존 extract_embedding이 파일 경로를 받으므로)
            temp_audio_path = self.audio_segmenter.save_temp_segment(audio_segment)
            
            # wav2vec2로 임베딩 추출
            embedding = extract_embedding(temp_audio_path)
            
            # MLP로 감정 예측
            with torch.no_grad():
                prediction = self.emotion_model(embedding)
                probabilities = torch.softmax(prediction, dim=-1).squeeze().tolist()
            
            # 임시 파일 정리
            self.audio_segmenter.cleanup_temp_file(temp_audio_path)
            
            return probabilities
            
        except Exception as e:
            print(f"❌ 음성 감정 분석 오류: {e}")
            return [0.2] * 5  # 기본값 반환
    
    def print_emotion_result(self, text, emotion_probs):
        """감정 분석 결과 출력"""
        print(f"\n🎵 음성 감정 분석 결과:")
        for label, prob in zip(self.emotion_labels, emotion_probs):
            print(f"  {label}: {prob:.3f}")
        
        # 가장 높은 확률의 감정
        max_idx = np.argmax(emotion_probs)
        max_emotion = self.emotion_labels[max_idx]
        max_prob = emotion_probs[max_idx]
        
        print(f"🎯 예측 감정: {max_emotion} ({max_prob:.3f})")
        print("-" * 50)
    
    def get_latest_audio_emotion(self):
        """Late Fusion용 최신 음성 감정 결과 반환"""
        if not self.audio_emotion_queue.empty():
            return self.audio_emotion_queue.get()
        return None
    
    def stop_processing(self):
        """처리 중단"""
        self.is_processing = False
        self.stt.stop_listening()
        self.audio_segmenter.stop_recording()
        print("🛑 실시간 처리 중단")


if __name__ == "__main__":
    processor = RealTimeEmotionProcessor()
    
    try:
        processor.start_processing()
        
        # 메인 루프 - Ctrl+C로 종료
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n종료 중...")
        processor.stop_processing()