app.py

import gradio as gr
from PIL import Image
import numpy as np
import onnxruntime as rt
import os
import spaces
import torch
from transformers import AutoImageProcessor 
from scipy.special import softmax
import requests # <-- **NEW:** สำหรับการเรียก Hugging Face API

# 1. การกำหนดค่าและโหลดโมเดล ONNX และ LLM
# ----------------------------------------------------
ONNX_MODEL_PATH = "model.onnx"
CLASS_LABELS_FILE = "class_labels.txt"
MODEL_ID = 'facebook/convnext-tiny-224' # ID โมเดลที่ใช้ในการฝึก

# **LLM CONFIGURATION (NEW)**
MODEL_NAME = "scb10x/typhoon2.5-qwen3-4b" # <--- ID โมเดล Typhoon 2.5
# ดึง Hugging Face Token จาก Environment Variable (ต้องตั้งค่าใน Space Secrets)
HF_TOKEN = os.getenv("HF_TOKEN") 
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"


# โหลดชื่อคลาส (ตัวละครวันพีซ)
try:
    with open(CLASS_LABELS_FILE, 'r', encoding='utf-8') as f:
        CHARACTER_LABELS = [line.strip() for line in f.readlines()]
except FileNotFoundError:
    CHARACTER_LABELS = ['Luffy', 'Zoro', 'Nami', 'Sanji', 'Chopper', 'Franky', 'Brook', 'Usopp', 'Jinbei', 'Robin', 'Ace', 'Law', 'Shanks', 'Kurohige', 'Mihawk', 'Rayleigh']
    print(f"⚠️ WARNING: {CLASS_LABELS_FILE} not found. Using default labels.")

# โหลด ONNX Runtime Session
try:
    sess = rt.InferenceSession(ONNX_MODEL_PATH)
    onnx_input_name = sess.get_inputs()[0].name
    onnx_output_name = sess.get_outputs()[0].name
    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
    print("ONNX model and Image Processor loaded successfully.")
except Exception as e:
    print(f"Error loading ONNX model or Image Processor: {e}. ตรวจสอบไฟล์ model.onnx และชื่อไฟล์.")
    sess = None


# 2. ฟังก์ชันเรียกใช้ Typhoon 2.5 API (แทนฟังก์ชันจำลองเดิม)
# ----------------------------------------------------
def query_typhoon_api(payload):
    """ส่ง Prompt ไปยัง Hugging Face Inference API ของ Typhoon 2.5"""
    if not HF_TOKEN:
        return "Error: HF_TOKEN is not set in Hugging Face Space secrets."
    
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    
    # ตรวจสอบ Response
    if response.status_code != 200:
        return f"Error {response.status_code}: API call failed. {response.text}"
        
    # ดึงผลลัพธ์
    try:
        # ผลลัพธ์จาก API จะอยู่ในรูปแบบ [ {'generated_text': '...'} ]
        result = response.json()[0]['generated_text']
        # โมเดลอาจจะตอบกลับด้วย Prompt เดิม เราจึงลบส่วน Prompt ออก
        return result.split(payload['inputs'])[-1].strip()
    except Exception as e:
        return f"Error processing API response: {e}"


# ฐานข้อมูลข้อมูลเสริมตัวละคร (ยังคงเก็บไว้เป็นความรู้พื้นฐาน)
CHARACTER_INFO = {
    # ... (ข้อมูลตัวละครยังคงเหมือนเดิม) ...
    "Ace": "โพโทกัส ดี เอส พี่ชายบุญธรรมของลูฟี่ ผู้ใช้พลังผลปีศาจเมระ เมระ",
    "Luffy": "มังกี้ ดี ลูฟี่ กัปตันกลุ่มโจรสลัดหมวกฟาง ผู้ใฝ่ฝันจะเป็นราชาโจรสลัด",
    "Zoro": "โรโรโนอา โซโล นักดาบสามเล่มแห่งกลุ่มหมวกฟาง ผู้มีเป้าหมายเป็นนักดาบอันดับหนึ่งของโลก",
    "Nami": "นามิ นักเดินเรือสาวแห่งกลุ่มหมวกฟาง และเป็นนักทำแผนที่มือฉมัง",
    "Sanji": "ซันจิ กุ๊กแห่งกลุ่มโจรสลัดหมวกฟาง และเป็นสุดยอดนักสู้ที่ใช้เท้าในการต่อสู้",
    "Chopper": "โทนี่ โทนี่ ช็อปเปอร์ หมอประจำเรือ ผู้มีใจรักเพื่อนและอ่อนไหวที่สุดในกลุ่ม",
    "Robin": "นิโค โรบิน นักโบราณคดี ผู้เดียวที่อ่านโพเนกลีฟได้",
    "Usopp": "อุซป พลซุ่มยิงและนักประดิษฐ์ ผู้มีความฝันเป็นนักรบผู้กล้าหาญแห่งท้องทะเล",
    "Franky": "แฟรงกี้ ช่างต่อเรือผู้สร้างเรือเธาซันด์ ซันนี่ มีพลังไซบอร์กสุดแกร่ง",
    "Brook": "บรู๊ค นักดนตรีผู้ใช้ดาบและมีชีวิตเป็นโครงกระดูก ผู้รักการร้องเพลงและมุกตลก",
    "Jinbei": "จินเบ อดีตเจ็ดเทพโจรสลัด และเป็นมนุษย์เงือกผู้เชี่ยวชาญคาราเต้เงือก",
    "Law": "ทราฟาลการ์ ลอว์ กัปตันกลุ่มโจรสลัดฮาร์ท ผู้ใช้พลังผลโอเปะ โอเปะ",
    "Shanks": "แชงค์ส หนึ่งในสี่จักรพรรดิ์ ผู้มอบหมวกฟางให้กับลูฟี่",
    "Mihawk": "จูราคิล มิฮอว์ค สุดยอดนักดาบผู้เป็นที่มาของฉายา 'ตาเหยี่ยว'",
    "Kurohige": "มาร์แชล ดี ทีช หรือหนวดดำ ผู้เป็นหนึ่งในสี่จักรพรรดิ์คนปัจจุบัน",
    "Rayleigh": "ซิลเวอร์ส เรย์ลี่ อดีตมือขวาของราชาโจรสลัด โกลด์ ดี. โรเจอร์",
}


def generate_thai_response(character_name, confidence):
    """
    ฟังก์ชัน LLM ที่ใช้ API จริง
    """
    info = CHARACTER_INFO.get(character_name, "ตัวละครวันพีซ")
    
    # 1. สร้าง Prompt ที่ชัดเจนสำหรับ Typhoon 2.5 
    prompt = (
        f"จากผลการวิเคราะห์ภาพ (ความมั่นใจ {confidence*100:.2f}%), ตัวละครที่ทำนายคือ '{character_name}'. "
        f"ตัวละครนี้คือ {info}. "
        f"กรุณาสร้างข้อความตอบกลับที่เป็นมิตรและเป็นภาษาไทย โดยขึ้นต้นด้วย 'ยืนยันผลการทำนาย!' "
        f"และรวมข้อมูลทั้งหมดนี้เข้าด้วยกันในประโยคเดียวโดยใช้ Markdown bold สำหรับชื่อตัวละครและความมั่นใจ (XX.XX%)."
    )
    
    # 2. เตรียม Payload สำหรับ API
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": 100,
            "return_full_text": False
        }
    }
    
    # 3. เรียก API และรับคำตอบ
    llm_response = query_typhoon_api(payload)
    
    # 4. หาก API call สำเร็จ ให้คืนคำตอบ หากล้มเหลวให้คืน Error
    if llm_response.startswith("Error"):
        return f"⚠️ LLM Error: {llm_response}"

    # Typhoon 2.5 มักจะตอบกลับเป็นประโยคที่ถูกต้องตาม Prompt
    return llm_response

# 3. ฟังก์ชันทำนายชื่อตัวละครด้วย ONNX (รองรับ ZeroGPU)
@spaces.GPU
def predict_one_piece_character(pil_image):
    if pil_image is None or sess is None:
        return "⚠️ โมเดลไม่พร้อมใช้งาน กรุณาตรวจสอบไฟล์ ONNX และการตั้งค่า"
    
    try:
        # 3.1 เตรียม Input ด้วย Image Processor
        inputs = processor(images=pil_image, return_tensors="np")
        onnx_input = inputs['pixel_values'].astype(np.float32)

        # 3.2 ทำนายผลด้วย ONNX Runtime
        onnx_predictions = sess.run([onnx_output_name], {onnx_input_name: onnx_input})
        logits = onnx_predictions[0].squeeze() 

        # 3.3 ประมวลผลผลลัพธ์ (Softmax และ Argmax)
        probabilities = softmax(logits)
        predicted_index = np.argmax(probabilities)
        predicted_character = CHARACTER_LABELS[predicted_index]
        confidence = probabilities[predicted_index].item()
        
        # 3.4 การรวมผล: ใช้ LLM Logic สร้างข้อความตอบกลับ
        final_response = generate_thai_response(predicted_character, confidence)
        
        return final_response

    except Exception as e:
        print(f"ERROR during prediction: {e}") 
        return f"เกิดข้อผิดพลาดในการทำนาย: {e}"

# 4. การสร้าง Gradio Interface
interface = gr.Interface(
    fn=predict_one_piece_character, 
    inputs=gr.Image(type="pil", label="อัปโหลดรูปภาพตัวละครวันพีซ"), 
    outputs=gr.Textbox(label="ผลการทำนายชื่อตัวละคร (Powered by Typhoon 2.5)"), 
    title="🏴‍☠️ One Piece Classifier (ConvNeXt ONNX + Typhoon 2.5)",
    description="อัปโหลดภาพตัวละครวันพีซ เพื่อให้ AI ทำนายชื่อตัวละครพร้อมสร้างข้อความตอบกลับที่น่าประทับใจ"
)

if __name__ == "__main__":
    interface.launch(inbrowser=True)