Update app.py

app.py

@@ -1,15 +1,15 @@
 import gradio as gr
-import openai, os, io, tempfile
+import openai, os, io, tempfile, mimetypes
 from dotenv import load_dotenv
 
-# ============== 환경 변수 & OpenAI 초기화 ===========================
+# ============ 환경 변수 & OpenAI 초기화 ==============================
 load_dotenv()
 api_key = os.getenv("OPENAI_API_KEY")
 if not api_key:
     raise RuntimeError("OPENAI_API_KEY를 .env 파일에 설정하세요!")
 client = openai.OpenAI(api_key=api_key)
 
-# ============== 언어 설정 ===========================================
+# ============ 언어 매핑 ==============================================
 LANGUAGES = [
     "Korean", "English", "Japanese", "Chinese",
     "Thai", "Russian", "Vietnamese",
@@ -20,33 +20,34 @@ LANG_CODE = {
     "Thai": "th", "Russian": "ru", "Vietnamese": "vi",
     "Spanish": "es", "French": "fr"
 }
-VOICE = {
-    lang: ("nova" if lang in ["Korean", "Japanese", "Chinese"] else "alloy")
-    for lang in LANGUAGES
-}
+VOICE = {l: ("nova" if l in ["Korean", "Japanese", "Chinese"] else "alloy")
+         for l in LANGUAGES}
+
+# ============ 공통 유틸 함수 =========================================
+def _safe_path(x):
+    """Gradio File/Audio 컴포넌트 값에서 실제 파일 경로 추출"""
+    if x is None:
+        return None
+    if isinstance(x, dict):
+        return x.get("name")  # 4.x 형식
+    return x
 
-# ============== 공통 함수 ===========================================
 def _gpt_translate(text: str, src: str, tgt: str) -> str:
-    """GPT-3.5 번역"""
+    """GPT-3.5-Turbo 번역"""
     rsp = client.chat.completions.create(
         model="gpt-3.5-turbo",
         messages=[
-            {
-                "role": "system",
-                "content": (
-                    f"You are a professional translator. Translate the following {src} text to {tgt}. "
-                    f"Only provide the translation without additional commentary."
-                )
-            },
+            {"role": "system",
+             "content": f"You are a professional translator. Translate the following {src} text to {tgt}. "
+                        f"Only provide the translated text."},
             {"role": "user", "content": text}
         ],
-        temperature=0.3,
-        max_tokens=2048
+        temperature=0.3, max_tokens=4096
     )
     return rsp.choices[0].message.content.strip()
 
 def _tts(text: str, lang: str) -> str:
-    """TTS-1 음성 합성 → 임시 mp3 경로 반환"""
+    """OpenAI TTS-1 → mp3 파일 경로"""
     out = client.audio.speech.create(
         model="tts-1",
         voice=VOICE.get(lang, "alloy"),
@@ -57,9 +58,13 @@ def _tts(text: str, lang: str) -> str:
     tmp.close()
     return tmp.name
 
-def translate_audio(audio_path, src, tgt):
-    """단일 음성 파일 번역(TTS 포함)"""
-    with open(audio_path, "rb") as f:
+# ============ ① 음성 파일/마이크 번역 ================================
+def translate_audio(audio_in, src, tgt):
+    path = _safe_path(audio_in)
+    if not path or not os.path.exists(path):
+        return "⚠️ 음성 파일을 녹음하거나 업로드하세요.", "", None
+
+    with open(path, "rb") as f:
         stt = client.audio.transcriptions.create(
             model="whisper-1",
             file=f,
@@ -73,53 +78,82 @@ def translate_audio(audio_path, src, tgt):
     tts_path = _tts(translated, tgt)
     return original, translated, tts_path
 
-# ============== 실시간 스트리밍 전사/번역 ============================
-STREAM_CHUNK_SEC = 4  # Whisper 호출 주기(초)
-
-def stream_generator(mic_stream, src, tgt):
-    """마이크 스트림 -> 주기적 청크 번역(누적 출력)"""
+# ============ ② PDF / 이미지 번역 ===================================
+def translate_document(file_in, src, tgt):
+    path = _safe_path(file_in)
+    if not path or not os.path.exists(path):
+        return "⚠️ PDF 또는 이미지를 업로드하세요.", ""
+
+    text = ""
+    ext = os.path.splitext(path)[1].lower()
+    mime = mimetypes.guess_type(path)[0] or ""
+
+    try:
+        if ext == ".pdf" or "pdf" in mime:
+            import pdfplumber
+            with pdfplumber.open(path) as pdf:
+                pages = pdf.pages[:5]  # 데모: 최대 5쪽
+                text = "\n".join(p.extract_text() or "" for p in pages)
+        elif ext in [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".gif"] or "image" in mime:
+            from PIL import Image
+            import pytesseract
+            text = pytesseract.image_to_string(Image.open(path))
+        else:
+            return "⚠️ 지원하지 않는 파일 형식입니다.", ""
+    except Exception as e:
+        return f"❌ 텍스트 추출 실패: {type(e).__name__}: {e}", ""
+
+    text = text.strip()
+    if not text:
+        return "⚠️ 텍스트가 추출되지 않았습니다.", ""
+
+    translated = _gpt_translate(text, src, tgt)
+    return text, translated
+
+# ============ ③ 실시간 스트리밍 전사/번역 ===========================
+STREAM_CHUNK_SEC = 4
+
+def stream_generator(mic_stream, src, tgt="English"):
     buffer = io.BytesIO()
     wav_header = None
-    original_acc, translated_acc = "", ""
+    orig_acc, trans_acc = "", ""
 
     while True:
         chunk = mic_stream.recv()
-        if chunk is None:                 # 스트림 종료
+        if chunk is None:
             break
-
         if not wav_header:
-            wav_header = chunk[:44]       # WAV 헤더(16kHz 16-bit mono)
+            wav_header = chunk[:44]  # WAV 헤더
         buffer.write(chunk)
 
-        # 지정 시간만큼 쌓이면 Whisper 호출
         if buffer.getbuffer().nbytes > 16000 * 2 * STREAM_CHUNK_SEC:
             wav_bytes = wav_header + buffer.getvalue()
             with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                 tmp.write(wav_bytes)
                 tmp.close()
                 o, t, _ = translate_audio(tmp.name, src, tgt)
+            orig_acc += " " + o
+            trans_acc += " " + t
+            yield orig_acc.strip(), trans_acc.strip()
+            buffer = io.BytesIO()
 
-            original_acc += " " + o
-            translated_acc += " " + t
-            yield original_acc.strip(), translated_acc.strip()
-            buffer = io.BytesIO()         # 버퍼 초기화
-
-    # 마지막 남은 버퍼 처리
     if buffer.getbuffer().nbytes:
         wav_bytes = wav_header + buffer.getvalue()
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
             tmp.write(wav_bytes)
             tmp.close()
             o, t, _ = translate_audio(tmp.name, src, tgt)
+        yield (orig_acc + " " + o).strip(), (trans_acc + " " + t).strip()
 
-        yield (original_acc + " " + o).strip(), (translated_acc + " " + t).strip()
-
-# ============== 4개국 동시 번역 =====================================
+# ============ ④ 4개국 동시 번역 ====================================
 FOUR_LANGS = ["English", "Chinese", "Thai", "Russian"]
 
-def translate_audio_four(audio_path, src):
-    """음성 파일 → 원문 + 4개 언어 동시 번역"""
-    with open(audio_path, "rb") as f:
+def translate_audio_four(audio_in, src):
+    path = _safe_path(audio_in)
+    if not path or not os.path.exists(path):
+        return ["⚠️ 음성 파일을 녹음하거나 업로드하세요."] + [""] * 4
+
+    with open(path, "rb") as f:
         stt = client.audio.transcriptions.create(
             model="whisper-1",
             file=f,
@@ -132,13 +166,13 @@ def translate_audio_four(audio_path, src):
     outs = [original]
     for lang in FOUR_LANGS:
         outs.append(_gpt_translate(original, src, lang))
-    return outs  # 총 5개(원문+4언어)
+    return outs
 
-# ============== Gradio UI ===========================================
+# ============ Gradio UI =============================================
 with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
     with gr.Tabs():
-        # ① 마이크/파일 번역
-        with gr.TabItem("🎙️ 마이크/파일 번역"):
+        # 1) 마이크/파일(오디오) 번역
+        with gr.TabItem("🎙️ 마이크/오디오 번역"):
             src1 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
             tgt1 = gr.Dropdown(LANGUAGES, value="English", label="출력")
             mic1 = gr.Audio(
@@ -157,29 +191,26 @@ with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
                 outputs=[stt1, tlt1, out1]
             )
 
-        # ② 파일 전용 번역
-        with gr.TabItem("🎧 파일 전용 번역"):
+        # 2) PDF / 이미지 번역
+        with gr.TabItem("📄 문서/이미지 번역"):
             src2 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
             tgt2 = gr.Dropdown(LANGUAGES, value="English", label="출력")
-            file2 = gr.Audio(
-                sources=["upload"],
-                type="filepath",
-                label="오디오 파일 업로드"
+            file2 = gr.File(
+                label="PDF 또는 이미지 업로드",
+                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".gif"]
             )
             btn2 = gr.Button("번역")
-            stt2 = gr.Textbox(label="원문", lines=5)
-            tlt2 = gr.Textbox(label="번역", lines=5)
-            out2 = gr.Audio(label="TTS", type="filepath", autoplay=True)
+            original_doc = gr.Textbox(label="추출된 원문", lines=15)
+            translated_doc = gr.Textbox(label="번역 결과", lines=15)
 
             btn2.click(
-                translate_audio,
+                translate_document,
                 inputs=[file2, src2, tgt2],
-                outputs=[stt2, tlt2, out2]
+                outputs=[original_doc, translated_doc]
             )
 
-        # ③ 실시간 스트리밍 번역(Beta)
+        # 3) 실시간 번역(Beta)
         with gr.TabItem("⏱️ 실시간 번역 (Beta)"):
-            gr.Markdown("마이크를 켜면 3~4초 간격으로 자막이 갱신됩니다.")
             src3 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
             tgt3 = gr.Dropdown(LANGUAGES, value="English", label="출력")
             mic3 = gr.Audio(
@@ -195,9 +226,8 @@ with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
 
             mic3.stream(gen, inputs=[src3, tgt3], outputs=[stt3, tlt3])
 
-        # ④ 4개 언어 동시 번역
+        # 4) 4개 언어 동시 번역
         with gr.TabItem("🌏 4개 언어 동시"):
-            gr.Markdown("입력 음성을 **English / Chinese(简体) / Thai / Russian** 4개 언어로 동시에 번역합니다.")
             src4 = gr.Dropdown(LANGUAGES, value="Korean", label="입력 언어")
             aud4 = gr.Audio(
                 sources=["microphone", "upload"],
@@ -219,6 +249,6 @@ with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
                 outputs=[org4, en4, zh4, th4, ru4]
             )
 
-# ============== 앱 실행 =============================================
+# ============ 실행 ===================================================
 if __name__ == "__main__":
     app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)