app.py

import gradio as gr
import base64
import json
import pandas as pd
from openai import OpenAI
import tempfile
import os
import io



UPSTAGE_API_KEY = ""

# 이미지 파일을 base64 문자열로 인코딩하는 함수
def encode_img_to_base64(img_path):
    with open(img_path, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode("utf-8")

# 중첩된 JSON을 평탄화하여 테이블로 변환할 수 있도록 처리하는 함수
def flatten_to_rows(data):
    flat = {}
    for key, value in data.items():
        if isinstance(value, list):
            for i, item in enumerate(value):
                if isinstance(item, dict):
                    for sub_key, sub_val in item.items():
                        flat[f"{key}[{i}].{sub_key}"] = sub_val
                else:
                    flat[f"{key}[{i}]"] = item
        elif isinstance(value, dict):
            for sub_key, sub_val in value.items():
                flat[f"{key}.{sub_key}"] = sub_val
        else:
            flat[key] = value
    return flat

# 명함 이미지를 받아 정보 추출 및 CSV 저장을 처리하는 핵심 함수
def extract_info(image_path, UPSTAGE_API_KEY):

    
    # Upstage API 클라이언트 설정 (스키마 생성용)
    schema_client = OpenAI(
        api_key=UPSTAGE_API_KEY,
        base_url="https://api.upstage.ai/v1/information-extraction/schema-generation"
    )
    
    # Upstage API 클라이언트 설정 (정보 추출용)
    extraction_client = OpenAI(
        api_key=UPSTAGE_API_KEY,
        base_url="https://api.upstage.ai/v1/information-extraction"
    )

    base64_data = encode_img_to_base64(image_path)

    # 스키마 생성
    schema_response = schema_client.chat.completions.create(
        model="information-extract",
        messages=[{
            "role": "user",
            "content": [{
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{base64_data}"}
            }]
        }]
    )
    schema = json.loads(schema_response.choices[0].message.content)

    # 정보 추출
    extraction_response = extraction_client.chat.completions.create(
        model="information-extract",
        messages=[{
            "role": "user",
            "content": [{
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{base64_data}"}
            }]
        }],
        response_format={
            "type": "json_schema",
            "json_schema": {
                "name": "business_card_schema",
                "schema": schema["json_schema"]["schema"]
            }
        }
    )

    result = json.loads(extraction_response.choices[0].message.content)
    json_result = json.dumps(result, indent=2, ensure_ascii=False)

    # 테이블화
    flat_dict = flatten_to_rows(result)
    df_result = pd.DataFrame([flat_dict]).T.reset_index()
    df_result.columns = ["Field", "Value"]


    return json_result, df_result



# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📇 명함 정보 자동 추출기")
    gr.Markdown("명함 이미지를 업로드하면 정보를 추출하고, 다운로드 버튼을 눌러 CSV로 저장할 수 있습니다.")

    with gr.Row():
        api_key_input = gr.Textbox(label="🔑 Upstage API Key", type="password", placeholder="Paste your API key here")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="명함 이미지 업로드")
            extract_btn = gr.Button("📤 정보 추출하기")

        with gr.Column():
            json_output = gr.Textbox(label="📄 추출된 JSON", lines=20)
            table_output = gr.Dataframe(label="📋 추출된 정보 (표 형태)", interactive=False) 

    extract_btn.click(
        fn=extract_info,
        inputs=[image_input, api_key_input],
        outputs=[json_output, table_output]
    )



# 앱 실행
if __name__ == "__main__":
    demo.launch()