import gradio as gr
from transformers import pipeline
from PIL import Image
import torch

# -----------------------------------------------------------------
# 1. 載入模型 (這會在 App 啟動時執行一次)
# (此部分與你的程式碼完全相同)
# -----------------------------------------------------------------
print("正在載入模型...")
# 模型 1: Image Captioning (圖片標註)
captioner = pipeline(
    "image-to-text",
    model="Salesforce/blip-image-captioning-base")

# 模型 2: Visual Question Answering (VQA)
vqa = pipeline(
    "visual-question-answering",
    model="dandelin/vilt-b32-finetuned-vqa")
print("模型載入完成！")

# -----------------------------------------------------------------
# 2. 定義 AI 功能的函式
# (此部分與你的程式碼完全相同)
# -----------------------------------------------------------------
def get_image_caption(pil_image):
    """
    接收一個 PIL 圖片，返回 AI 生成的文字描述。
    """
    if pil_image is None:
        return "錯誤：請先上傳一張圖片。"
    try:
        caption = captioner(pil_image)
        if caption and isinstance(caption, list) and 'generated_text' in caption[0]:
            return caption[0]['generated_text'].capitalize() # 讓首字母大寫
        else:
            return "錯誤：模型返回了非預期的格式。"
    except Exception as e:
        return f"處理時發生錯誤: {e}"

def answer_visual_question(pil_image, question_text):
    """
    接收一個 PIL 圖片和一個問題字串，返回 AI 的回答。
    """
    if pil_image is None:
        return "錯誤：請先上傳一張圖片。"
    if not question_text:
        return "錯誤：請輸入你的問題。"
    try:
        answer = vqa(image=pil_image, question=question_text)
        if answer and isinstance(answer, list) and 'answer' in answer[0]:
            return answer[0]['answer'].capitalize() # 讓首字母大寫
        else:
            return "錯誤：模型返回了非預期的格式。"
    except Exception as e:
        return f"處理時發生錯誤: {e}"

# -----------------------------------------------------------------
# 3. (全新設計) 建立 Gradio Web UI
# -----------------------------------------------------------------

# (新增) 自訂 CSS 樣式
# 我們將使用這些 CSS 來打造卡片佈局和美化元件
my_custom_css = """
/* --- 1. 整體頁面背景 --- */
.gradio-container {
    background: linear-gradient(to right, #e0eafc, #cfdef3); /* 淡雅的漸層背景 */
    padding: 20px;
}

/* --- 2. 美化主標題 --- */
#main-title {
    text-align: center;
    color: #2c3e50;
    font-size: 3em;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    margin-bottom: 20px;
}

/* --- 3. 卡片式佈局 (gr.Group) --- */
.gradio-group {
    background: rgba(255, 255, 255, 0.9); /* 半透明白色背景 */
    border-radius: 20px;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.15); /* 更柔和的陰影 */
    border: 1px solid rgba(255, 255, 255, 0.18);
    padding: 25px !important;
    margin-bottom: 20px !important;
}

/* --- 4. 美化功能標題 (h3) --- */
.gradio-group h3 {
    color: #34495e;
    font-size: 1.5em;
    border-bottom: 2px solid #4a90e2; /* 標題底線 */
    padding-bottom: 5px;
}

/* --- 5. 美化主要按鈕 --- */
.gr-button-primary {
    background: linear-gradient(to right, #4a90e2, #50E3C2) !important;
    border: none !important;
    border-radius: 10px !important;
    font-size: 1.1em !important;
    color: white !important;
    padding: 12px 24px !important;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important;
    transition: all 0.3s ease !important;
}
.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 12px rgba(0,0,0,0.15) !important;
}

/* --- 6. 美化輸出文字框 (使其看起來像"顯示面板") --- */
#output-caption textarea, #output-vqa textarea {
    background: #f4f7f6 !important;
    border: 2px solid #bdc3c7 !important;
    border-radius: 10px !important;
    font-size: 1.15em !important;
    color: #2c3e50 !important;
    font-weight: 500 !important;
}
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=my_custom_css) as demo:

    # App 標題 (使用 elem_id 以便 CSS 選取)
    gr.Markdown(
        """
        # 🎨 視覺語言 AI 展示 (Vision Language AI Demo)
        """,
        elem_id="main-title"
    )
    gr.Markdown("### 統一上傳圖片，然後在下方體驗兩種不同的 AI 功能。")

    with gr.Row(variant="panel"):
        # --- (重新設計) 佈局：左側為輸入 ---
        with gr.Column(scale=2):
            with gr.Group():
                gr.Markdown("### 1. 上傳圖片 (Input)")
                # 建立一個共用的圖片輸入元件
                shared_image_input = gr.Image(type="pil", label="上傳圖片")
                
                gr.Markdown("#### 試試看範例圖片：")
                # 範例功能 (會自動填入上面的 shared_image_input)
                gr.Examples(
                    examples=[
                        ["sample_images/cat.jpg"],
                        ["sample_images/dog.jpg"]
                    ],
                    inputs=shared_image_input,
                    label="點選範例圖片"
                )

        # --- (重新設計) 佈局：右側為功能輸出 ---
        with gr.Column(scale=3):
            # --- 功能一：圖片標註 ---
            with gr.Group():
                gr.Markdown("### 2. 功能一：圖片標註 (Image Captioning)")
                caption_button = gr.Button("生成圖片描述", variant="primary")
                caption_output_text = gr.Textbox(
                    label="AI 生成的描述", 
                    interactive=False, 
                    elem_id="output-caption" # 綁定 CSS ID
                )
            
            # --- 功能二：視覺問答 (VQA) ---
            with gr.Group():
                gr.Markdown("### 3. 功能二：視覺問答 (VQA)")
                vqa_input_question = gr.Textbox(
                    label="你的問題", 
                    placeholder="例如：圖片中有幾個人？ (What color is the car?)"
                )
                vqa_button = gr.Button("對圖片提出問題", variant="primary")
                vqa_output_answer = gr.Textbox(
                    label="AI 的回答", 
                    interactive=False, 
                    elem_id="output-vqa" # 綁定 CSS ID
                )
                
                # VQA 的範例 (會同時填入圖片 和 問題)
                gr.Examples(
                    examples=[
                        ["sample_images/cat.jpg", "What color is the cat?"],
                        ["sample_images/dog.jpg", "What color is the dog?"]
                    ],
                    inputs=[shared_image_input, vqa_input_question],
                    label="點選範例 (包含圖片與問題)"
                )

    # -----------------------------------------------------------------
    # 4. 綁定按鈕和函式 (邏輯不變，只是元件名稱更新)
    # -----------------------------------------------------------------
    
    # 功能一：
    caption_button.click(
        fn=get_image_caption,
        inputs=[shared_image_input],  # 輸入來自共用的圖片
        outputs=[caption_output_text]
    )

    # 功能二：
    vqa_button.click(
        fn=answer_visual_question,
        inputs=[shared_image_input, vqa_input_question], # 輸入來自共用圖片 + VQA問題
        outputs=[vqa_output_answer]
    )

# -----------------------------------------------------------------
# 5. 啟動 App (此部分與你的程式碼相同)
# -----------------------------------------------------------------
if __name__ == "__main__":
    demo.launch(debug=True)