import numpy as np
import gradio as gr
from PIL import Image
import joblib

# 加载保存的KNN模型
model = joblib.load('best_knn_model.pkl')

def preprocess_image(input_data):
    if isinstance(input_data, dict):
        image_data = input_data.get('composite', None)
        if image_data is None:
            image_data = input_data.get('image', None)
    else:
        image_data = input_data

    if image_data is None or np.all(image_data == 0):
        print("Empty or invalid image data received.")
        return None

    pil_image = Image.fromarray(image_data).convert('L')
    pil_image = pil_image.resize((8, 8), Image.LANCZOS)
    image_array = np.array(pil_image).astype(float)  # Ensure it's floating point for further processing
    image_array = (image_array / 255.0 * 16.0)  # Normalize and scale
    return image_array.flatten()

# 定义预测函数
def predict_digit(image):
    if image is None:
        return {}, ""
    
    # 预处理图像
    processed_image = preprocess_image(image)
    
    # 进行预测
    prediction = model.predict([processed_image])[0]
    probabilities = model.predict_proba([processed_image])[0]
    
    # 创建结果字典
    results = {str(i): float(prob) for i, prob in enumerate(probabilities)}
    
    return results, str(prediction)

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="Draw a digit here", mode="L"),  # 修改为 mode="L"
    outputs=[gr.Label(num_top_classes=3), gr.Textbox(label="预测结果")],
    live=False,  # 关闭实时模式
    title="手写数字识别",
    description="画一个数字（0-9），然后点击确定按钮，模型将尝试识别它。",
    allow_flagging="never",  # 禁用标记功能
)

# 启动Gradio接口
iface.launch(share=True)