app.py

import gradio as gr
from PIL import Image
import numpy as np
import onnxruntime as rt
import os
import spaces
import torch
# Use AutoImageProcessor for robust image preprocessing (ConvNeXt/ViT)
from transformers import AutoImageProcessor 
from scipy.special import softmax # Used for computing probabilities from model output
import requests # Used for making HTTP API calls (Typhoon 2.5)

# 1. ONNX MODEL AND LLM CONFIGURATION
# ----------------------------------------------------
ONNX_MODEL_PATH = "model.onnx"
CLASS_LABELS_FILE = "class_labels.txt"
MODEL_ID = 'facebook/convnext-tiny-224' # The base model ID used for training

# **LLM CONFIGURATION (Actual API Call Setup)**
MODEL_NAME = "scb10x/typhoon2.5-qwen3-4b" # Typhoon 2.5 Model ID
# Fetch Hugging Face Token from Space Secrets
HF_TOKEN = os.getenv("HF_TOKEN") 
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"


# Load character classes from the file created during training
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.")

# Load ONNX Runtime Session
try:
    print(f"Attempting to load ONNX model from: {ONNX_MODEL_PATH}") # <-- NEW DEBUG LINE
    sess = rt.InferenceSession(ONNX_MODEL_PATH)
    onnx_input_name = sess.get_inputs()[0].name
    onnx_output_name = sess.get_outputs()[0].name
    # Load Image Processor (essential for correct image preparation)
    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
    print("ONNX model and Image Processor loaded successfully.")
except Exception as e:
    # <--- NEW DEBUG BLOCK: พิมพ์ Error จริงออกมา
    print(f"FATAL ERROR LOADING ONNX MODEL: {e}")
    print("Please ensure model.onnx is tracked by Git LFS and is uploaded correctly.")
    sess = None


# 2. LLM API FUNCTION (Replaces Placeholder)
# ----------------------------------------------------
def query_typhoon_api(payload):
    """Sends prompt to the Hugging Face Inference API of 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)
    
    # Check for non-successful status codes (e.g., 401 Unauthorized, 503 Service Unavailable)
    if response.status_code != 200:
        return f"Error {response.status_code}: API call failed. {response.text}"
        
    try:
        # Extract the generated text from the response structure
        result = response.json()[0]['generated_text']
        # Remove the input prompt part from the output text
        return result.split(payload['inputs'])[-1].strip()
    except Exception as e:
        return f"Error processing API response: {e}"


# 3. TYPHOON 2.5 LOGIC (Knowledge Base + Prompt Generation)
# ---------------------------------------------------------
# This dictionary serves as the LLM's knowledge base for the character's full name/role
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):
    """
    Constructs a sophisticated prompt and queries the Typhoon 2.5 API.
    """
    info = CHARACTER_INFO.get(character_name, "ตัวละครวันพีซ")
    
    # 1. Build a clear, instructional prompt for the LLM
    prompt = (
        f"จากผลการวิเคราะห์ภาพ (ความมั่นใจ {confidence*100:.2f}%), ตัวละครที่ทำนายคือ '{character_name}'. "
        f"ตัวละครนี้คือ {info}. "
        f"กรุณาสร้างข้อความตอบกลับที่เป็นมิตรและเป็นภาษาไทย โดยขึ้นต้นด้วย 'ยืนยันผลการทำนาย!' "
        f"และรวมข้อมูลทั้งหมดนี้เข้าด้วยกันในประโยคเดียวโดยใช้ Markdown bold สำหรับชื่อตัวละครและความมั่นใจ (XX.XX%)."
    )
    
    # 2. Prepare Payload
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": 100,
            "return_full_text": False
        }
    }
    
    # 3. Call the API and handle potential errors
    llm_response = query_typhoon_api(payload)
    
    if llm_response.startswith("Error"):
        # Fallback to a static, simple response if API fails
        thai_name = info.split(' ')[0] 
        return (f"⚠️ LLM API ไม่ตอบสนอง: ตัวละครคือ **{thai_name}** ({info}) "
                f"[ความมั่นใจ: **{confidence*100:.2f}%**]")

    return llm_response

# 4. ONNX INFERENCE FUNCTION
# ----------------------------------------------------
#@spaces.GPU
def predict_one_piece_character(pil_image):
    if pil_image is None or sess is None:
        return "⚠️ โมเดลไม่พร้อมใช้งาน กรุณาตรวจสอบไฟล์ ONNX และการตั้งค่า"
    
    try:
        # 4.1 Preprocessing (ConvNeXt standard input)
        inputs = processor(images=pil_image, return_tensors="np")
        onnx_input = inputs['pixel_values'].astype(np.float32)

        # 4.2 Run Inference
        onnx_predictions = sess.run([onnx_output_name], {onnx_input_name: onnx_input})
        logits = onnx_predictions[0].squeeze() 

        # 4.3 Post-processing (Softmax and Argmax)
        probabilities = softmax(logits)
        predicted_index = np.argmax(probabilities)
        predicted_character = CHARACTER_LABELS[predicted_index]
        confidence = probabilities[predicted_index].item()
        
        # 4.4 LLM Integration
        final_response = generate_thai_response(predicted_character, confidence)
        
        return final_response

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

# 5. 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)