Update app.py

app.py

@@ -15,10 +15,11 @@ import torch
 import numpy as np
 from loguru import logger
 from PIL import Image
-from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer, WhisperProcessor, WhisperForConditionalGeneration
 import time
 import warnings
 from typing import Dict, List, Optional, Union
+import librosa
 
 # CSV/TXT 분석
 import pandas as pd
@@ -27,7 +28,7 @@ import PyPDF2
 
 warnings.filterwarnings('ignore')
 
-print("🎮 로봇 시각 시스템 초기화 (Gemma3-R1984-4B)...")
+print("🎮 로봇 시각 시스템 초기화 (Gemma3-R1984-4B + Whisper)...")
 
 ##############################################################################
 # 상수 정의
@@ -42,7 +43,10 @@ SERPHOUSE_API_KEY = os.getenv("SERPHOUSE_API_KEY", "")
 ##############################################################################
 model = None
 processor = None
+whisper_model = None
+whisper_processor = None
 model_loaded = False
+whisper_loaded = False
 model_name = "Gemma3-R1984-4B"
 
 ##############################################################################
@@ -54,6 +58,72 @@ def clear_cuda_cache():
         torch.cuda.empty_cache()
         gc.collect()
 
+##############################################################################
+# Whisper 모델 로드
+##############################################################################
+@spaces.GPU(duration=60)
+def load_whisper():
+    global whisper_model, whisper_processor, whisper_loaded
+    
+    if whisper_loaded:
+        logger.info("Whisper 모델이 이미 로드되어 있습니다.")
+        return True
+    
+    try:
+        logger.info("Whisper 모델 로딩 시작...")
+        whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+        whisper_model = WhisperForConditionalGeneration.from_pretrained(
+            "openai/whisper-base",
+            device_map="auto",
+            torch_dtype=torch.float16
+        )
+        whisper_loaded = True
+        logger.info("✅ Whisper 모델 로딩 완료!")
+        return True
+    except Exception as e:
+        logger.error(f"Whisper 모델 로딩 실패: {e}")
+        return False
+
+##############################################################################
+# 오디오 처리 함수
+##############################################################################
+@spaces.GPU(duration=30)
+def transcribe_audio(audio_data):
+    """Whisper를 사용한 오디오 전사"""
+    global whisper_model, whisper_processor
+    
+    if not whisper_loaded:
+        if not load_whisper():
+            return "오디오 처리 불가"
+    
+    try:
+        if audio_data is None:
+            return None
+        
+        # 오디오 데이터 처리
+        sample_rate, audio = audio_data
+        
+        # 16kHz로 리샘플링
+        if sample_rate != 16000:
+            audio = librosa.resample(audio.astype(float), orig_sr=sample_rate, target_sr=16000)
+        
+        # Whisper 입력 처리
+        inputs = whisper_processor(audio, sampling_rate=16000, return_tensors="pt")
+        inputs = {k: v.to(whisper_model.device) for k, v in inputs.items()}
+        
+        # 음성 인식
+        with torch.no_grad():
+            generated_ids = whisper_model.generate(**inputs, max_length=225)
+        
+        # 디코딩
+        transcription = whisper_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        return transcription.strip()
+        
+    except Exception as e:
+        logger.error(f"오디오 전사 오류: {e}")
+        return f"오디오 인식 실패: {str(e)}"
+
 ##############################################################################
 # 키워드 추출 함수
 ##############################################################################
@@ -85,8 +155,8 @@ def do_web_search(query: str) -> str:
             "domain": "google.com",
             "serp_type": "web",
             "device": "desktop",
-            "lang": "ko",  # 한국어 우선
-            "num": "10"   # 10개로 제한
+            "lang": "ko",
+            "num": "10"
         }
         
         headers = {
@@ -232,10 +302,11 @@ def analyze_image_for_robot(
     prompt: str,
     task_type: str = "general",
     use_web_search: bool = False,
-    enable_thinking: bool = False,  # 기본값 False로 변경
-    max_new_tokens: int = 300  # 장면 설명을 위해 300으로 증가
+    enable_thinking: bool = False,
+    max_new_tokens: int = 300,
+    audio_transcript: Optional[str] = None
 ) -> str:
-    """로봇 작업을 위한 이미지 분석"""
+    """로봇 작업을 위한 이미지 분석 (오디오 정보 포함)"""
     global model, processor
     
     if not model_loaded:
@@ -247,7 +318,7 @@ def analyze_image_for_robot(
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image).convert('RGB')
         
-        # 태스크별 시스템 프롬프트 구성 (더 간결하게)
+        # ���스크별 시스템 프롬프트 구성
         system_prompts = {
             "general": "당신은 로봇 시각 시스템입니다. 먼저 장면을 1-2줄로 설명하고, 핵심 내용을 간결하게 분석하세요.",
             "planning": """당신은 로봇 작업 계획 AI입니다. 
@@ -265,6 +336,18 @@ Step_n: xxx""",
             "pointing": "당신은 지점 지정 시스템입니다. 먼저 참조점들을 한 줄로 설명하고, 위치를 [(x1,y1), (x2,y2), ...]로 반환하세요."
         }
         
+        # 오디오 정보가 있으면 프롬프트 수정
+        if audio_transcript and task_type == "planning":
+            system_prompts["planning"] = """당신은 로봇 작업 계획 AI입니다. 
+먼저 장면 이해를 1-2줄로 설명하고, 주변 소리를 인식했다면 그것도 설명한 후, 작업 계획을 작성하세요.
+형식: 
+[장면 이해] 현재 보이는 장면을 1-2줄로 설명
+[주변 소리 인식] 들리는 소리나 음성을 1줄로 설명
+[작업 계획]
+Step_1: xxx
+Step_2: xxx
+Step_n: xxx"""
+        
         system_prompt = system_prompts.get(task_type, system_prompts["general"])
         
         # Chain-of-Thought 추가 (선택적)
@@ -280,6 +363,11 @@ Step_n: xxx""",
                 search_results = do_web_search(keywords)
                 combined_system = f"{search_results}\n\n{system_prompt}"
         
+        # 사용자 프롬프트에 오디오 정보 추가
+        user_prompt = prompt
+        if audio_transcript:
+            user_prompt += f"\n\n[인식된 주변 소리: {audio_transcript}]"
+        
         # 메시지 구성
         messages = [
             {
@@ -290,7 +378,7 @@ Step_n: xxx""",
                 "role": "user",
                 "content": [
                     {"type": "image", "url": image},
-                    {"type": "text", "text": prompt}
+                    {"type": "text", "text": user_prompt}
                 ]
             }
         ]
@@ -494,30 +582,23 @@ css = """
     background: #e8f5e9;
     color: #2e7d32;
 }
+.audio-status {
+    text-align: center;
+    padding: 5px;
+    border-radius: 5px;
+    margin: 5px 0;
+    font-weight: bold;
+    background: #e3f2fd;
+    color: #1565c0;
+}
 """
 
 with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as demo:
     gr.HTML("""
     <div class="robot-header">
         <h1>🤖 로봇 시각 시스템</h1>
-        <h3>🎮 Gemma3-R1984-4B + 📷 실시간 웹캠 + 🔍 웹 검색</h3>
-        <p>⚡ 최신 멀티모달 AI로 로봇 작업 분석 및 계획 수립!</p>
-    </div>
-    """)
-    
-    gr.HTML("""
-    <div class="info-box">
-        <h4>🌟 시스템 특징:</h4>
-        <ul>
-            <li>🖼️ 고급 이미지/비디오 분석 (Gemma3-4B VLM)</li>
-            <li>👁️ 장면 이해 및 상황 설명</li>
-            <li>📋 다단계 작업 계획 및 추론</li>
-            <li>📍 정밀한 객체 위치 파악 (Grounding)</li>
-            <li>🤏 로봇 파지점 분석 (Affordance)</li>
-            <li>🛤️ 경로 계획 (Trajectory Planning)</li>
-            <li>🔍 실시간 웹 검색 통합</li>
-            <li>🔄 10초마다 자동 캡처 및 분석</li>
-        </ul>
+        <h3>🎮 Gemma3-R1984-4B + 📷 실시간 웹캠 + 🎤 음성 인식</h3>
+        <p>⚡ 멀티모달 AI로 로봇 작업 분석!</p>
     </div>
     """)
     
@@ -532,7 +613,7 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
                     streaming=True,
                     type="numpy",
                     label="실시간 스트리밍",
-                    height=350
+                    height=300
                 )
                 
                 # 자동 캡처 상태 표시
@@ -543,30 +624,47 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
                 # 캡처된 이미지 표시
                 captured_image = gr.Image(
                     label="캡처된 이미지",
-                    height=200,
+                    height=180,
                     visible=False
                 )
             
+            # 오디오 컨트롤
+            gr.Markdown("### 🎤 음성 인식")
+            with gr.Group():
+                # 오디오 상태 표시
+                audio_status = gr.HTML(
+                    '<div class="audio-status">🎤 음성 인식: 비활성화</div>'
+                )
+                
+                # 마지막 인식된 텍스트
+                last_transcript = gr.Textbox(
+                    label="인식된 음성",
+                    value="",
+                    lines=2,
+                    interactive=False
+                )
+            
             # 로봇 작업 버튼들
-            gr.Markdown("### 🎯 로봇 작업 선택")
+            gr.Markdown("### 🎯 로봇 작업")
             with gr.Row():
                 capture_btn = gr.Button("📸 수동 캡처", variant="primary", elem_classes="task-button")
                 clear_capture_btn = gr.Button("🗑️ 초기화", elem_classes="task-button")
             
-            with gr.Row():
+            with gr.Column():
                 auto_capture_toggle = gr.Checkbox(
-                    label="🔄 자동 캡처 활성화 (10초마다)",
+                    label="🔄 자동 캡처 (10초마다)",
+                    value=False
+                )
+                
+                use_audio_toggle = gr.Checkbox(
+                    label="🎤 음성 인식 사용",
                     value=False,
-                    info="활성화 시 10초마다 자동으로 캡처 및 분석"
+                    info="주변 소리를 인식하여 분석에 포함"
                 )
             
             with gr.Row():
                 planning_btn = gr.Button("📋 작업 계획", elem_classes="task-button")
                 grounding_btn = gr.Button("📍 객체 위치", elem_classes="task-button")
-            
-            with gr.Row():
-                affordance_btn = gr.Button("🤏 파지점 분석", elem_classes="task-button")
-                trajectory_btn = gr.Button("🛤️ 경로 계획", elem_classes="task-button")
         
         # 오른쪽: 분석 설정 및 결과
         with gr.Column(scale=2):
@@ -575,7 +673,7 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
             with gr.Row():
                 with gr.Column():
                     task_prompt = gr.Textbox(
-                        label="작업 설명 / 질문",
+                        label="작업 설명",
                         placeholder="예: 테이블 위의 컵을 잡아서 싱크대에 놓기",
                         value="현재 장면을 분석하고 로봇이 수행할 수 있는 작업을 제안하세요.",
                         lines=2
@@ -583,40 +681,46 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
                     
                     with gr.Row():
                         use_web_search = gr.Checkbox(
-                            label="🔍 웹 검색 사용",
-                            value=False,
-                            info="관련 정보를 웹에서 검색합니다"
+                            label="🔍 웹 검색",
+                            value=False
                         )
                         
                         enable_thinking = gr.Checkbox(
-                            label="🤔 추론 과정 표시",
-                            value=False,  # 기본값 False로 변경
-                            info="Chain-of-Thought 추론 과정을 보여줍니다"
+                            label="🤔 추론 과정",
+                            value=False
                         )
                     
                     max_tokens = gr.Slider(
-                        label="최대 토큰 수",
+                        label="최대 토큰",
                         minimum=100,
-                        maximum=4096,
-                        value=300,  # 장면 설명을 위해 300으로 증가
+                        maximum=1000,
+                        value=300,
                         step=50
                     )
             
             gr.Markdown("### 📊 분석 결과")
             result_output = gr.Textbox(
                 label="AI 분석 결과",
-                lines=20,
-                max_lines=40,
+                lines=18,
+                max_lines=35,
                 show_copy_button=True,
                 elem_id="result"
             )
             
             status_display = gr.HTML(
-                '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비 완료</div>'
+                '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비</div>'
             )
     
-    # 문서 분석 탭 (숨김 처리)
-    with gr.Tab("📄 문서 분석", visible=False):  # visible=False로 숨김
+    # 숨겨진 오디오 입력
+    audio_input = gr.Audio(
+        sources=["microphone"],
+        streaming=True,
+        visible=False,
+        label="마이크 입력"
+    )
+    
+    # 문서 분석 탭 (숨김)
+    with gr.Tab("📄 문서 분석", visible=False):
         with gr.Row():
             with gr.Column():
                 doc_files = gr.File(
@@ -648,7 +752,8 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
     
     # 이벤트 핸들러
     webcam_state = gr.State(None)
-    auto_capture_state = gr.State({"enabled": False, "timer": None})
+    audio_state = gr.State(None)
+    transcript_state = gr.State("")
     
     def capture_webcam(frame):
         """웹캠 프레임 캡처"""
@@ -658,9 +763,9 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
     
     def clear_capture():
         """캡처 초기화"""
-        return None, gr.update(visible=False), '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비 완료</div>'
+        return None, gr.update(visible=False), '<div class="status-box" style="background:#d4edda; color:#155724;">🎮 시스템 준비</div>', ""
     
-    def analyze_with_task(image, prompt, task_type, use_search, thinking, tokens):
+    def analyze_with_task(image, prompt, task_type, use_search, thinking, tokens, transcript):
         """특정 태스크로 이미지 분석"""
         if image is None:
             return "❌ 먼저 이미지를 캡처하세요.", '<div class="status-box" style="background:#f8d7da; color:#721c24;">❌ 이미지 없음</div>'
@@ -673,10 +778,11 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
             task_type=task_type,
             use_web_search=use_search,
             enable_thinking=thinking,
-            max_new_tokens=tokens
+            max_new_tokens=tokens,
+            audio_transcript=transcript if transcript else None
         )
         
-        # 결과 포맷팅 (더 간결하게)
+        # 결과 포맷팅
         timestamp = time.strftime("%H:%M:%S")
         task_names = {
             "planning": "작업 계획",
@@ -694,19 +800,28 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
         return formatted_result, complete_status
     
     # 자동 캡처 및 분석 함수
-    def auto_capture_and_analyze(webcam_frame, task_prompt, use_search, thinking, tokens, auto_state):
-        """자동 캡처 및 분석"""
+    def auto_capture_and_analyze(webcam_frame, audio_data, task_prompt, use_search, thinking, tokens, use_audio, current_transcript):
+        """자동 캡처 및 분석 (오디오 포함)"""
         if webcam_frame is None:
             return (
                 None,
                 "자동 캡처 대기 중...",
                 '<div class="status-box" style="background:#fff3cd; color:#856404;">⏳ 웹캠 대기 중</div>',
-                '<div class="auto-capture-status">🔄 자동 캡처: 웹캠 대기 중</div>'
+                '<div class="auto-capture-status">🔄 자동 캡처: 웹캠 대기 중</div>',
+                current_transcript,
+                current_transcript
             )
         
         # 캡처 수행
         timestamp = time.strftime("%H:%M:%S")
         
+        # 오디오 처리 (활성화된 경우)
+        new_transcript = ""
+        if use_audio and audio_data is not None:
+            transcribed = transcribe_audio(audio_data)
+            if transcribed and transcribed != "오디오 처리 불가":
+                new_transcript = transcribed
+        
         # 이미지 분석 (작업 계획 모드로)
         result = analyze_image_for_robot(
             image=webcam_frame,
@@ -714,7 +829,8 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
             task_type="planning",
             use_web_search=use_search,
             enable_thinking=thinking,
-            max_new_tokens=tokens
+            max_new_tokens=tokens,
+            audio_transcript=new_transcript if new_transcript else None
         )
         
         formatted_result = f"""🔄 자동 분석 완료 ({timestamp})
@@ -726,7 +842,9 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
             webcam_frame,
             formatted_result,
             '<div class="status-box" style="background:#d4edda; color:#155724;">✅ 자동 분석 완료</div>',
-            f'<div class="auto-capture-status">🔄 자동 캡처: 마지막 분석 {timestamp}</div>'
+            f'<div class="auto-capture-status">🔄 자동 캡처: 마지막 분석 {timestamp}</div>',
+            new_transcript if new_transcript else current_transcript,
+            new_transcript if new_transcript else current_transcript
         )
     
     # 웹캠 스트리밍
@@ -736,6 +854,16 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
         outputs=[webcam_state]
     )
     
+    # 오디오 스트리밍
+    def process_audio_stream(audio_data):
+        return audio_data
+    
+    audio_input.stream(
+        fn=process_audio_stream,
+        inputs=[audio_input],
+        outputs=[audio_state]
+    )
+    
     # 수동 캡처 버튼
     capture_btn.click(
         fn=capture_webcam,
@@ -746,31 +874,19 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
     # 초기화 버튼
     clear_capture_btn.click(
         fn=clear_capture,
-        outputs=[webcam_state, captured_image, status_display]
+        outputs=[webcam_state, captured_image, status_display, transcript_state]
     )
     
     # 작업 버튼들
     planning_btn.click(
-        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "planning", s, t, tk),
-        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
+        fn=lambda img, p, s, t, tk, tr: analyze_with_task(img, p, "planning", s, t, tk, tr),
+        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens, transcript_state],
         outputs=[result_output, status_display]
     )
     
     grounding_btn.click(
-        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "grounding", s, t, tk),
-        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
-        outputs=[result_output, status_display]
-    )
-    
-    affordance_btn.click(
-        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "affordance", s, t, tk),
-        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
-        outputs=[result_output, status_display]
-    )
-    
-    trajectory_btn.click(
-        fn=lambda img, p, s, t, tk: analyze_with_task(img, p, "trajectory", s, t, tk),
-        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens],
+        fn=lambda img, p, s, t, tk, tr: analyze_with_task(img, p, "grounding", s, t, tk, tr),
+        inputs=[captured_image, task_prompt, use_web_search, enable_thinking, max_tokens, transcript_state],
         outputs=[result_output, status_display]
     )
     
@@ -791,7 +907,7 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
     )
     
     # 자동 캡처 타이머 (10초마다)
-    timer = gr.Timer(10.0, active=False)  # 10초 타이머, 초기에는 비활성화
+    timer = gr.Timer(10.0, active=False)
     
     # 자동 캡처 토글 이벤트
     def toggle_auto_capture(enabled):
@@ -806,11 +922,32 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
         outputs=[timer, auto_capture_status]
     )
     
+    # 오디오 토글 이벤트
+    def toggle_audio(enabled):
+        if enabled:
+            # Whisper 모델 로드
+            load_whisper()
+            return (
+                gr.update(visible=True),  # audio_input 표시
+                '<div class="audio-status">🎤 음성 인식: 활성화됨</div>'
+            )
+        else:
+            return (
+                gr.update(visible=False),  # audio_input 숨김
+                '<div class="audio-status">🎤 음성 인식: 비활성화</div>'
+            )
+    
+    use_audio_toggle.change(
+        fn=toggle_audio,
+        inputs=[use_audio_toggle],
+        outputs=[audio_input, audio_status]
+    )
+    
     # 타이머 틱 이벤트
     timer.tick(
         fn=auto_capture_and_analyze,
-        inputs=[webcam_state, task_prompt, use_web_search, enable_thinking, max_tokens, auto_capture_state],
-        outputs=[captured_image, result_output, status_display, auto_capture_status]
+        inputs=[webcam_state, audio_state, task_prompt, use_web_search, enable_thinking, max_tokens, use_audio_toggle, transcript_state],
+        outputs=[captured_image, result_output, status_display, auto_capture_status, transcript_state, last_transcript]
     )
     
     # 초기 모델 로드
@@ -824,7 +961,7 @@ with gr.Blocks(title="🤖 로봇 시각 시스템 (Gemma3-4B)", css=css) as dem
     )
 
 if __name__ == "__main__":
-    print("🚀 로봇 시각 시스템 시작 (Gemma3-R1984-4B)...")
+    print("🚀 로봇 시각 시스템 시작 (Gemma3-R1984-4B + Whisper)...")
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,