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

# 1. 載入模型 (這會自動下載模型，第一次執行會比較久)
# ---------------------------------------------------
# 載入 Image Captioning 模型
captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")

# 載入 VQA 模型
vqa = pipeline("visual-question-answering", model="dandelin/vilt-b32-finetuned-vqa")


# 2. 定義 AI 功能的函式
# ---------------------------------------------------
def image_captioning(image):
    """
    看圖說故事
    """
    if image is None:
        return "請先上傳圖片"
    # pipeline 會回傳一個 list，裡面是 dict
    result = captioner(image)
    return result[0]['generated_text']

def visual_question_answer(image, question):
    """
    視覺問答
    """
    if image is None:
        return "請先上傳圖片"
    if not question:
        return "請輸入問題"
        
    result = vqa(image, question)
    # VQA 的回傳格式不同，直接取 'answer'
    return result[0]['answer']

# 3. 建立 Gradio 介面
# ---------------------------------------------------
with gr.Blocks(title="Vision Language AI 應用") as demo:
    gr.Markdown("# 🎨 Vision Language AI 應用展示")
    gr.Markdown("這是一個展示 Vision-Language (VL) 模型的 Web UI，由 Gradio 和 Hugging Face Transformers 驅動。")

    with gr.Tabs():
        
        # --- 分頁 1: Image Captioning ---
        with gr.TabItem("🖼️ 看圖說故事 (Image Captioning)"):
            with gr.Row():
                image_input_caption = gr.Image(type="pil", label="上傳圖片")
                text_output_caption = gr.Textbox(label="AI 產生的描述")
            caption_button = gr.Button("開始產生描述")

        # --- 分頁 2: Visual Question Answering ---
        with gr.TabItem("❓ 視覺問答 (VQA)"):
            with gr.Row():
                image_input_vqa = gr.Image(type="pil", label="上傳圖片")
                with gr.Column():
                    question_input_vqa = gr.Textbox(label="請輸入你對圖片的問題")
                    text_output_vqa = gr.Textbox(label="AI 的回答")
            vqa_button = gr.Button("開始問答")

    # 4. 綁定按鈕和函式
    # ---------------------------------------------------
    
    # 綁定 "看圖說故事" 的按鈕
    caption_button.click(
        fn=image_captioning,
        inputs=[image_input_caption],
        outputs=[text_output_caption]
    )
    
    # 綁定 "視覺問答" 的按鈕
    vqa_button.click(
        fn=visual_question_answer,
        inputs=[image_input_vqa, question_input_vqa],
        outputs=[text_output_vqa]
    )
    
    # 增加一些範例
    gr.Markdown("--- \n ### 範例圖片")
    gr.Examples(
        examples=[
            ["sample_images/cat_dog.jpg"],
            ["sample_images/bus.jpg"]
        ],
        inputs=[image_input_caption],
        outputs=[text_output_caption],
        fn=image_captioning,
        cache_examples=True # 範例可以先跑過
    )
    gr.Examples(
        examples=[
            ["sample_images/bus.jpg", "What color is the bus?"],
            ["sample_images/cat_dog.jpg", "How many animals are in the picture?"]
        ],
        inputs=[image_input_vqa, question_input_vqa],
        outputs=[text_output_vqa],
        fn=visual_question_answer,
        cache_examples=True
    )


# 5. 啟動 Web App
# ---------------------------------------------------
if __name__ == "__main__":
    demo.launch() # share=True 會產生一個公開網址，但作業要求用 HF Spaces，所以本地跑不用開