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

# --- 1. 載入 AI 模型 ---
# 使用 Hugging Face pipeline 快速載入預訓練模型
# 任務: "image-to-text" (圖片轉文字)
# 模型: "nlpconnect/vit-gpt2-image-captioning"
#
# 第一次執行時，會自動下載模型 (可能需要幾分鐘)
print("正在載入 AI 模型...")
# 確保模型在 CPU 上運行 (符合作業的 Zero GPU/CPU 要求)
captioner = pipeline("image-to-text", 
                     model="nlpconnect/vit-gpt2-image-captioning",
                     device="cpu") 
print("模型載入完成。")

# --- 2. 定義核心 AI 功能 ---
def generate_caption(input_image: Image.Image) -> str:
    """
    接收一個 PIL Image 物件，返回 AI 生成的描述文字。
    
    Args:
        input_image: Gradio 傳入的 PIL 圖片物件。

    Returns:
        一個描述圖片的字串。
    """
    if input_image is None:
        return "錯誤：請先上傳圖片。"
        
    try:
        # 使用 pipeline 進行預測
        # pipeline 會返回一個 list，其中包含一個 dict
        # result 格式: [{'generated_text': 'a cat sitting on a couch'}]
        result = captioner(input_image)
        caption = result[0]['generated_text']
        
        # 對輸出做一點清理，讓它更易讀
        return caption.capitalize()
        
    except Exception as e:
        print(f"處理圖片時發生錯誤: {e}")
        return f"發生錯誤: {e}"

# --- 3. 建立 Gradio Web UI 介面 ---

# 定義輸入介面：
# type="pil" 確保傳入 Python 函式的是 PIL Image 物件
input_component = gr.Image(type="pil", label="上傳圖片 (Upload Image)")

# 定義輸出介面：
output_component = gr.Textbox(label="AI 產生的圖片描述 (AI-Generated Caption)")

# 讀取範例圖片 (請確保你已在 sample_images/ 資料夾放入對應檔案)
# 這是為了讓使用者一點就知道怎麼用
example_list = [
    ["sample_images/cat.jpg"],
    ["sample_images/dog_park.jpg"],
]

# 組合介面
iface = gr.Interface(
    fn=generate_caption,       # 要呼叫的函式
    inputs=input_component,    # 輸入元件
    outputs=output_component,  # 輸出元件
    
    # 介面標題與說明
    title="Vision Language AI：智慧圖片描述",
    description="這是一個使用 Gradio 和 Hugging Face Transformers 打造的 Web App。\n上傳一張圖片，AI 將會自動為您產生一段描述文字。",
    
    # 提供範例
    examples=example_list,
    
    # 介面主題 (可選)
    theme=gr.themes.Soft()
)

# --- 4. 啟動 Web App ---
if __name__ == "__main__":
    # share=True 會產生一個公開的臨時網址，方便分享
    # (但在 HF Spaces 上部署時不需要 share=True)
    iface.launch()