app.py

import gradio as gr
import base64
import requests
import json
from PIL import Image
import io
import os
from typing import Optional, Tuple
import re

class KoreanOCRApp:
    def __init__(self):
        self.api_key = None
        self.project_id = None
        
    def set_credentials(self, api_key: str, project_id: str) -> str:
        """API 키와 프로젝트 ID 설정 및 검증"""
        if not api_key or not project_id:
            return "❌ API 키와 프로젝트 ID를 모두 입력해주세요."
        
        # 프로젝트 ID 검증 (영문, 숫자, 하이픈만 허용)
        if not re.match(r'^[a-z0-9\-]+$', project_id.strip()):
            return "❌ 유효하지 않은 프로젝트 ID 형식입니다. 영문 소문자, 숫자, 하이픈만 사용 가능합니다."
        
        self.api_key = api_key.strip()
        self.project_id = project_id.strip()
        return "✅ 인증 정보가 설정되었습니다."
    
    def encode_image_to_base64(self, image: Image.Image) -> str:
        """이미지를 base64로 인코딩"""
        buffer = io.BytesIO()
        # JPEG 형식으로 저장하여 파일 크기 최적화
        if image.mode == 'RGBA':
            # RGBA 이미지는 RGB로 변환
            background = Image.new('RGB', image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[-1])
            image = background
        
        image.save(buffer, format='JPEG', quality=95)
        image_bytes = buffer.getvalue()
        return base64.b64encode(image_bytes).decode('utf-8')
    
    def call_gemini_api_direct(self, image_base64: str) -> str:
        """Gemini API 직접 호출 (Google AI Studio API 사용)"""
        if not self.api_key:
            return "❌ 먼저 API 키를 설정해주세요."
        
        # Google AI Studio API 엔드포인트 사용
        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:generateContent?key={self.api_key}"
        
        headers = {
            "Content-Type": "application/json"
        }
        
        payload = {
            "contents": [{
                "parts": [
                    {
                        "text": """이 이미지에 포함된 모든 한국어 텍스트를 정확하게 추출해주세요. 
                        다음 규칙을 따라주세요:
                        1. 이미지에서 발견되는 모든 한국어 텍스트를 순서대로 추출
                        2. 텍스트의 위치나 레이아웃을 최대한 보존
                        3. 줄바꿈과 문단 구분을 명확히 표시
                        4. 특수문자, 숫자, 영어가 포함되어 있다면 그대로 유지
                        5. 읽기 어려운 부분이 있다면 [불분명] 표시
                        
                        추출된 텍스트만 반환해주세요."""
                    },
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": image_base64
                        }
                    }
                ]
            }],
            "generationConfig": {
                "temperature": 0.1,
                "topP": 0.8,
                "topK": 40,
                "maxOutputTokens": 8192
            },
            "safetySettings": [
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH", 
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                }
            ]
        }
        
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=60)
            
            if response.status_code == 401:
                return "❌ API 키가 유효하지 않습니다. Google AI Studio에서 발급받은 올바른 API 키를 입력해주세요."
            elif response.status_code == 403:
                return "❌ API 접근 권한이 없습니다. Gemini API가 활성화되어 있는지 확인해주세요."
            elif response.status_code == 429:
                return "❌ API 호출 한도를 초과했습니다. 잠시 후 다시 시도해주세요."
            
            response.raise_for_status()
            
            result = response.json()
            
            if "candidates" in result and len(result["candidates"]) > 0:
                content = result["candidates"][0]["content"]["parts"][0]["text"]
                return content.strip()
            elif "error" in result:
                return f"❌ API 오류: {result['error'].get('message', '알 수 없는 오류')}"
            else:
                return "❌ 텍스트를 추출할 수 없습니다. 이미지에 한국어 텍스트가 포함되어 있는지 확인해주세요."
                
        except requests.exceptions.RequestException as e:
            return f"❌ API 호출 오류: {str(e)}"
        except json.JSONDecodeError:
            return "❌ API 응답 파싱 오류가 발생했습니다."
        except KeyError as e:
            return f"❌ 예상치 못한 API 응답 형식: {str(e)}"
        except Exception as e:
            return f"❌ 알 수 없는 오류: {str(e)}"
    
    def call_vertex_ai_api(self, image_base64: str) -> str:
        """Vertex AI API 호출 (서비스 계정 키 사용)"""
        if not self.api_key or not self.project_id:
            return "❌ 먼저 API 키와 프로젝트 ID를 설정해주세요."
        
        location = "us-central1"
        url = f"https://{location}-aiplatform.googleapis.com/v1/projects/{self.project_id}/locations/{location}/publishers/google/models/gemini-1.5-pro:generateContent"
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "contents": [{
                "role": "user",
                "parts": [
                    {
                        "text": """이 이미지에 포함된 모든 한국어 텍스트를 정확하게 추출해주세요. 
                        다음 규칙을 따라주세요:
                        1. 이미지에서 발견되는 모든 한국어 텍스트를 순서대로 추출
                        2. 텍스트의 위치나 레이아웃을 최대한 보존
                        3. 줄바꿈과 문단 구분을 명확히 표시
                        4. 특수문자, 숫자, 영어가 포함되어 있다면 그대로 유지
                        5. 읽기 어려운 부분이 있다면 [불분명] 표시
                        
                        추출된 텍스트만 반환해주세요."""
                    },
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": image_base64
                        }
                    }
                ]
            }],
            "generation_config": {
                "temperature": 0.1,
                "top_p": 0.8,
                "top_k": 40,
                "max_output_tokens": 8192
            }
        }
        
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=60)
            
            if response.status_code == 401:
                return "❌ 인증 오류: Access Token이 유효하지 않거나 만료되었습니다."
            elif response.status_code == 403:
                return "❌ 권한 오류: Vertex AI API 접근 권한이 없습니다."
            elif response.status_code == 404:
                return "❌ 프로젝트 ID가 올바르지 않거나 Vertex AI API가 활성화되지 않았습니다."
            
            response.raise_for_status()
            
            result = response.json()
            
            if "candidates" in result and len(result["candidates"]) > 0:
                content = result["candidates"][0]["content"]["parts"][0]["text"]
                return content.strip()
            else:
                return "❌ 텍스트를 추출할 수 없습니다. 이미지에 한국어 텍스트가 포함되어 있는지 확인해주세요."
                
        except requests.exceptions.RequestException as e:
            return f"❌ API 호출 오류: {str(e)}"
        except Exception as e:
            return f"❌ 알 수 없는 오류: {str(e)}"
    
    def process_image(self, image: Optional[Image.Image], api_key: str, project_id: str, api_type: str) -> Tuple[Optional[Image.Image], str]:
        """이미지 처리 및 OCR 수행"""
        if image is None:
            return None, "❌ 이미지를 업로드해주세요."
        
        # 인증 정보 설정
        if api_type == "Google AI Studio":
            if not api_key:
                return image, "❌ Google AI Studio API 키를 입력해주세요."
            self.api_key = api_key.strip()
        else:  # Vertex AI
            auth_result = self.set_credentials(api_key, project_id)
            if "❌" in auth_result:
                return image, auth_result
        
        try:
            # 이미지 크기 확인 및 조정
            img_byte_array = io.BytesIO()
            image.save(img_byte_array, format='JPEG', quality=95)
            img_size_mb = len(img_byte_array.getvalue()) / (1024 * 1024)
            
            if img_size_mb > 4:  # 4MB로 제한을 낮춤
                # 이미지 크기가 너무 크면 리사이즈
                max_dimension = 1920
                image.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
            
            # 이미지를 base64로 인코딩
            image_base64 = self.encode_image_to_base64(image)
            
            # API 타입에 따라 호출
            if api_type == "Google AI Studio":
                extracted_text = self.call_gemini_api_direct(image_base64)
            else:
                extracted_text = self.call_vertex_ai_api(image_base64)
            
            # 결과 반환
            return image, extracted_text
            
        except Exception as e:
            return image, f"❌ 이미지 처리 중 오류가 발생했습니다: {str(e)}"

# 전역 앱 인스턴스
ocr_app = KoreanOCRApp()

def create_interface():
    """Gradio 인터페이스 생성"""
    
    # CSS 스타일링
    css = """
    .gradio-container {
        font-family: 'Noto Sans KR', sans-serif;
    }
    
    .main-header {
        text-align: center;
        color: #2c3e50;
        margin-bottom: 20px;
    }
    
    .info-box {
        background-color: #e8f4fd;
        border: 1px solid #bee5eb;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    
    .warning-box {
        background-color: #fff3cd;
        border: 1px solid #ffeaa7;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
        color: #856404;
    }
    """
    
    with gr.Blocks(css=css, title="한국어 OCR - Gemini AI") as interface:
        
        gr.Markdown("""
        # 🔍 한국어 OCR 텍스트 추출기
        ### Google Gemini AI를 활용한 고정밀 한국어 문자 인식
        
        이미지에서 한국어 텍스트를 정확하게 추출합니다. 문서, 간판, 손글씨 등 다양한 형태의 한국어를 인식할 수 있습니다.
        """, elem_classes="main-header")
        
        # API 선택
        gr.Markdown("## 🔧 API 설정")
        
        api_type = gr.Radio(
            choices=["Google AI Studio", "Vertex AI"],
            value="Google AI Studio",
            label="사용할 API 선택",
            info="Google AI Studio는 개인 사용자용, Vertex AI는 기업용"
        )
        
        # 인증 정보 입력 섹션
        with gr.Row():
            with gr.Column(scale=2):
                api_key_input = gr.Textbox(
                    label="API 키 / Access Token",
                    placeholder="Google AI Studio API 키 또는 Vertex AI Access Token",
                    type="password",
                    lines=1
                )
            with gr.Column(scale=1):
                project_id_input = gr.Textbox(
                    label="프로젝트 ID (Vertex AI만)",
                    placeholder="Google Cloud 프로젝트 ID",
                    lines=1
                )
        
        # API 설정 가이드
        with gr.Accordion("📖 API 설정 가이드", open=False):
            gr.Markdown("""
            ### Google AI Studio API (권장)
            1. [Google AI Studio](https://aistudio.google.com/)에 접속
            2. "Get API Key" 클릭
            3. API 키 생성 및 복사
            4. 위의 "API 키" 필드에 붙여넣기
            
            ### Vertex AI API (고급 사용자용)
            1. [Google Cloud Console](https://console.cloud.google.com/)에서 프로젝트 생성
            2. Vertex AI API 활성화
            3. 서비스 계정 생성 및 키 다운로드
            4. `gcloud auth application-default login` 또는 Access Token 발급
            5. API 키와 프로젝트 ID 입력
            
            ### ⚠️ 주의사항
            - Google AI Studio는 개인 사용자에게 무료 할당량 제공
            - Vertex AI는 유료 서비스로 사용량에 따라 과금
            - API 키는 안전하게 보관하고 공유하지 마세요
            """, elem_classes="warning-box")
        
        # 이미지 업로드 및 처리 섹션
        gr.Markdown("## 📤 이미지 업로드 및 텍스트 추출")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="📁 이미지 업로드",
                    type="pil",
                    sources=["upload", "clipboard"],
                    interactive=True
                )
                
                process_btn = gr.Button(
                    "🔍 텍스트 추출 시작", 
                    variant="primary",
                    size="lg"
                )
                
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="📋 업로드된 이미지 확인",
                    type="pil",
                    interactive=False
                )
        
        # 추출된 텍스트 출력
        gr.Markdown("## 📝 추출된 텍스트")
        
        extracted_text = gr.Textbox(
            label="인식된 한국어 텍스트",
            placeholder="추출된 텍스트가 여기에 표시됩니다...",
            lines=10,
            max_lines=20,
            interactive=True,
            show_copy_button=True
        )
        
        # 이벤트 핸들러
        process_btn.click(
            fn=ocr_app.process_image,
            inputs=[input_image, api_key_input, project_id_input, api_type],
            outputs=[output_image, extracted_text],
            show_progress=True
        )
        
        # 사용 팁
        gr.Markdown("""
        ### 💡 사용 팁
        
        **📸 이미지 품질:**
        - 선명하고 해상도가 높은 이미지 사용
        - 충분한 조명과 대비
        - 텍스트가 수평으로 배치된 이미지 권장
        
        **📄 지원 형식:**
        - **이미지 형식:** PNG, JPEG, WebP
        - **최대 크기:** 4MB (자동 리사이즈)
        - **인식 언어:** 한국어, 영어, 숫자, 특수문자
        
        **🔒 보안:**
        - API 키는 세션 동안만 임시 저장
        - 이미지는 서버에 저장되지 않음
        - 개인정보가 포함된 이미지 사용 시 주의
        
        **⚡ 성능:**
        - Google AI Studio: 빠르고 안정적 (권장)
        - Vertex AI: 기업용 고급 기능
        """)
    
    return interface

# 메인 실행
if __name__ == "__main__":
    # 인터페이스 생성 및 실행
    demo = create_interface()
    
    # 서버 실행
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        debug=True,
        show_error=True,
        inbrowser=True
    )