app.py

import gradio as gr
import os
import logging
from datetime import datetime
from stt_processor import TextProcessor

# 로깅 설정
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 전역 변수
text_processor = None

def initialize_processor():
    """텍스트 프로세서를 초기화합니다."""
    global text_processor
    try:
        # 환경 변수 또는 Hugging Face Secrets에서 API 키 읽기
        google_api_key = os.getenv("GOOGLE_API_KEY")
        
        if not google_api_key:
            return False, "❌ Google API 키가 설정되지 않았습니다. Hugging Face Spaces의 Settings에서 GOOGLE_API_KEY를 설정해주세요."
        
        text_processor = TextProcessor(google_api_key)
        return True, "✅ 텍스트 프로세서가 초기화되었습니다."
        
    except Exception as e:
        logger.error(f"텍스트 프로세서 초기화 실패: {e}")
        return False, f"❌ 초기화 실패: {str(e)}"

def process_text_input(input_text, progress=gr.Progress()):
    """
    입력된 텍스트를 처리합니다.
    
    Args:
        input_text: 처리할 텍스트
        progress: Gradio 진행률 객체
        
    Returns:
        tuple: (처리 상태, 원본 텍스트, 화자 분리 결과, 교정 결과, 화자1 대화, 화자2 대화)
    """
    global text_processor
    
    if not input_text or not input_text.strip():
        return "❌ 처리할 텍스트를 입력해주세요.", "", "", "", "", ""
    
    try:
        # 텍스트 프로세서 초기화 (필요한 경우)
        if text_processor is None:
            progress(0.1, desc="텍스트 프로세서 초기화 중...")
            success, message = initialize_processor()
            if not success:
                return message, "", "", "", "", ""
        
        # 모델 로딩
        progress(0.2, desc="AI 모델 로딩 중...")
        if not text_processor.models_loaded:
            text_processor.load_models()
        
        # 진행 상황 콜백 함수
        def progress_callback(status, current, total):
            progress_value = 0.2 + (current / total) * 0.7  # 0.2~0.9 범위
            progress(progress_value, desc=f"{status} ({current}/{total})")
        
        # 텍스트 처리
        progress(0.3, desc="텍스트 처리 시작...")
        result = text_processor.process_text(input_text, progress_callback=progress_callback)
        
        if not result.get("success", False):
            return f"❌ 처리 실패: {result.get('error', 'Unknown error')}", "", "", "", "", ""
        
        # 결과 추출
        progress(0.95, desc="결과 정리 중...")
        original_text = result["original_text"]
        separated_text = result["separated_text"]
        corrected_text = result["corrected_text"]
        
        # 화자별 대화 추출
        conversations = result["conversations_by_speaker_corrected"]
        speaker1_text = "\n\n".join([f"{i+1}. {utterance}" for i, utterance in enumerate(conversations.get("화자1", []))])
        speaker2_text = "\n\n".join([f"{i+1}. {utterance}" for i, utterance in enumerate(conversations.get("화자2", []))])
        
        progress(1.0, desc="처리 완료!")
        
        status_message = f"""
✅ **처리 완료!**
- 텍스트명: {result['text_name']}
- 처리 시간: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
- 화자1 발언 수: {len(conversations.get('화자1', []))}개
- 화자2 발언 수: {len(conversations.get('화자2', []))}개
"""
        
        return status_message, original_text, separated_text, corrected_text, speaker1_text, speaker2_text
        
    except Exception as e:
        logger.error(f"텍스트 처리 중 오류: {e}")
        return f"❌ 처리 중 오류가 발생했습니다: {str(e)}", "", "", "", "", ""

def create_interface():
    """Gradio 인터페이스를 생성합니다."""
    
    # CSS 스타일링
    css = """
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .status-box {
        padding: 15px;
        border-radius: 8px;
        margin: 10px 0;
    }
    .main-header {
        text-align: center;
        color: #2c3e50;
        margin-bottom: 20px;
    }
    """
    
    with gr.Blocks(css=css, title="2인 대화 STT 처리기") as interface:
        
        # 헤더
        gr.HTML("""
        <div class="main-header">
            <h1>💬 2인 대화 화자 분리기 (AI)</h1>
            <p>Gemini 2.0 Flash AI를 사용한 텍스트 화자 분리 및 맞춤법 교정</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # 텍스트 입력 섹션
                gr.Markdown("### 📝 텍스트 입력")
                text_input = gr.Textbox(
                    label="2인 대화 텍스트를 입력하세요",
                    placeholder="두 명이 나누는 대화 내용을 여기에 붙여넣기하세요...\n\n예시:\n안녕하세요, 오늘 회의에 참석해주셔서 감사합니다. 네, 안녕하세요. 준비된 자료가 있나요? 네, 프레젠테이션 자료를 준비했습니다.",
                    lines=8,
                    max_lines=15
                )
                
                process_btn = gr.Button(
                    "🚀 처리 시작",
                    variant="primary",
                    size="lg"
                )
                
                # 상태 표시
                status_output = gr.Markdown(
                    "### 📊 처리 상태\n준비 완료. 2인 대화 텍스트를 입력하고 '처리 시작' 버튼을 클릭하세요.",
                    elem_classes=["status-box"]
                )
            
            with gr.Column(scale=2):
                # 결과 표시 섹션
                with gr.Tabs():
                    with gr.TabItem("📝 원본 텍스트"):
                        original_output = gr.Textbox(
                            label="입력된 원본 텍스트",
                            lines=10,
                            max_lines=20,
                            placeholder="처리 후 원본 텍스트가 여기에 표시됩니다..."
                        )
                    
                    with gr.TabItem("👥 화자 분리 (원본)"):
                        separated_output = gr.Textbox(
                            label="AI 화자 분리 결과 (원본)",
                            lines=10,
                            max_lines=20,
                            placeholder="처리 후 화자별로 분리된 대화가 여기에 표시됩니다..."
                        )
                    
                    with gr.TabItem("✏️ 화자 분리 (교정)"):
                        corrected_output = gr.Textbox(
                            label="AI 화자 분리 결과 (맞춤법 교정)",
                            lines=10,
                            max_lines=20,
                            placeholder="처리 후 맞춤법이 교정된 화자 분리 결과가 여기에 표시됩니다..."
                        )
                    
                    with gr.TabItem("👤 화자1 대화"):
                        speaker1_output = gr.Textbox(
                            label="화자1의 모든 발언",
                            lines=10,
                            max_lines=20,
                            placeholder="처리 후 화자1의 발언들이 여기에 표시됩니다..."
                        )
                    
                    with gr.TabItem("👤 화자2 대화"):
                        speaker2_output = gr.Textbox(
                            label="화자2의 모든 발언",
                            lines=10,
                            max_lines=20,
                            placeholder="처리 후 화자2의 발언들이 여기에 표시됩니다..."
                        )
        
        # 사용법 안내
        gr.Markdown("""
        ### 📖 사용법
        1. **텍스트 입력**: 2인 대화 텍스트를 입력란에 붙여넣기하세요
        2. **처리 시작**: '🚀 처리 시작' 버튼을 클릭하여 화자 분리를 시작하세요
        3. **결과 확인**: 각 탭에서 원본 텍스트, 화자 분리 결과, 개별 화자 대화를 확인하세요
        
        ### ⚙️ 기술 정보
        - **화자 분리**: Google Gemini 2.0 Flash
        - **맞춤법 교정**: 고급 AI 기반 한국어 교정
        - **지원 언어**: 한국어 최적화
        - **최적 환경**: 2인 대화, 명확한 문맥
        
        ### ⚠️ 주의사항
        - 처리 시간은 텍스트 길이에 따라 달라집니다 (보통 30초-2분)
        - Google AI API 사용량 제한이 있을 수 있습니다
        - 2인 대화에 최적화되어 있습니다
        - 대화 맥락이 명확할수록 분리 정확도가 높아집니다
        """)
        
        # 이벤트 연결
        process_btn.click(
            fn=process_text_input,
            inputs=[text_input],
            outputs=[
                status_output,
                original_output,
                separated_output,
                corrected_output,
                speaker1_output,
                speaker2_output
            ],
            show_progress=True
        )
    
    return interface

# 메인 실행
if __name__ == "__main__":
    logger.info("Gradio 앱을 시작합니다...")
    
    # 인터페이스 생성
    app = create_interface()
    
    # 앱 실행
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )