import os
import gradio as gr
import torch
import torchvision.models as models
from torch.serialization import safe_globals
from torchvision import transforms
import gradio_client.utils as client_utils

_orig = client_utils._json_schema_to_python_type
def _safe_json_schema_to_python_type(schema, defs=None):
    if isinstance(schema, bool):
        return "dict"  # 或者 "Any"
    return _orig(schema, defs)

client_utils._json_schema_to_python_type = _safe_json_schema_to_python_type



# 載入模型
with safe_globals([models.resnet.ResNet]):
    model = torch.load("model.pth", map_location="cpu", weights_only=False)
model.eval()

# 類別名稱
class_names = ["吉伊", "小八", "兔兔"]

# 圖片前處理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3),
])

# 推論函式
def classify_image(img):
    try:
        img = transform(img).unsqueeze(0)
        with torch.no_grad():
            outputs = model(img)
            probs = torch.nn.functional.softmax(outputs, dim=1)
        return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
    except Exception as e:
        return {"error": str(e)}

# Gradio 介面：用 JSON 輸出代替 Label
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="上傳圖片"),
    outputs=gr.Label(label="預測結果"),
    title="吉伊卡哇角色分類器",
    description="🐰 上傳吉伊、小八或兔兔的圖片，我來判斷是誰"
).queue()

if __name__ == "__main__":
    demo.launch(
        debug=True,
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860))
    )

if __name__ == "__main__":
    demo.launch(show_api=False, share=True)