# 导入必要的库和模块
import joblib  # 导入 joblib 库，用于加载预训练的模型
import gradio as gr  # 导入 Gradio 库，用于创建交互式界面
import numpy as np  # 导入 NumPy 库，用于数据处理
from PIL import Image  # 导入 PIL 库中的 Image 模块，用于图像处理

# 加载保存的 KNN 模型
knn_model = joblib.load(r'D:\gitee\optimal_knn_mnist_question\best_knn_model.pkl')

# # 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_knn(digit_draw):
    try:
        # 确保传入的 digit_draw 是一个 PIL 图像对象
        if isinstance(digit_draw, dict) and 'image' in digit_draw:
            image = digit_draw['image']  # 如果是字典类型，从中提取图像
        else:
            image = digit_draw  # 如果已经是图像对象，则直接使用

        # 将图像转换为灰度图
        image = image.convert('L')
        
        # 使用 LANCZOS 滤波器将图像大小调整为 8x8 像素
        image = image.resize((8, 8), Image.LANCZOS)
        
        # 将图像转换为数组并将其展平为一维数组
        image_array = np.array(image).flatten()
        
        # 归一化像素值，这里是除以 16，因为像素值范围是 0-255，而 KNN 模型可能在归一化的数据上训练
        image_array = image_array / 16
        
        # 使用 KNN 模型进行预测
        prediction = knn_model.predict([image_array])[0]
        
        # 返回预测结果
        return str(prediction)
    except Exception as e:
        # 如果发生异常，返回错误信息
        return f"Error: {str(e)}"

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
# 定义输入和输出
inputs = gr.Sketchpad(label="Draw your handwritten digit", type="pil")  # 创建一个画板输入组件
output = gr.Textbox(label="Predicted Digit")  # 创建一个文本框输出组件

# 创建Gradio界面
iface = gr.Interface(  # 创建 Gradio 接口
    fn=predict_knn,  # 将预测函数绑定到接口
    inputs=inputs,  # 设置输入组件
    outputs=output,  # 设置输出组件
    title='KNN Handwritten Digit Prediction',  # 设置接口标题
    description='Draw a digit in the sketchpad and submit to see the model prediction. Ensure the drawing is clear and centered.'  # 设置接口描述
)

# 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch()