from fastai.vision.all import *
import gradio as gr
import torch
import torch.nn.functional as F


learn = load_learner("export_new.pkl")
# 用模型的 vocab 生成中文标签
categories = tuple("数字" + str(x) for x in learn.dls.vocab)

def classify_image(img):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)
    if img.mode != 'L':
        img = img.convert('L') # 转换为灰度图
    img = img.resize((28, 28))
    img = PILImage.create(img)
    try:
        pred, idx, probs = learn.predict(img)
        print(f"预测类别: {pred}")
        print(f"类别索引: {idx}")
        print(f"概率分布: {probs}")
            # 手动归一化
        if not torch.is_tensor(probs):
            probs = torch.tensor(probs)
        probs = F.softmax(probs, dim=0)
        print("softmax后：", probs)
        return dict(zip(categories, map(float, probs)))
    except IndexError as e:
        print(f"错误: {e}")
        print(f"预测类别: {pred}", "检查模型的 vocab 是否匹配:", learn.dls.vocab)
        error_msg = f"预测失败: {str(e)}"
        raise RuntimeError(error_msg) from e

with gr.Blocks() as demo:
    gr.Markdown("# Digits Classifier")
    gr.Markdown("Classify digits numbers.")

    with gr.Row():
        image_input = gr.Image(height=192, width=192)
        label_output = gr.Label()

    btn = gr.Button("Classify")
    btn.click(fn=classify_image, inputs=image_input, outputs=label_output)

    gr.Examples(
        examples=[
            ["examples/number3.png"],
            ["examples/number6.png"],
            ["examples/number9.png"],
        ],
        inputs=image_input
    )

if __name__ == "__main__":
    demo.launch()