import gradio as gr
import pickle
import numpy as np
from PIL import Image
from sklearn.preprocessing import StandardScaler

# 加载保存的 KNN 模型
with open('best_knn_model.pkl', 'rb') as f:
    model = pickle.load(f)

# 标准化处理函数（根据训练时的处理）
scaler = StandardScaler()

def predict_digit(drawing):
    composite_image = drawing['composite']
    image = np.array(composite_image.convert('L')).astype(np.float32)
    image = Image.fromarray(image)
    image = image.resize((8, 8), Image.Resampling.LANCZOS)  # 使用正确的重采样方法
    image = np.array(image).flatten()  # 展平图像
    image = (255 - image) / 255.0
    image = scaler.fit_transform(image.reshape(1, -1))

    # 预测
    prediction = model.predict(image)
    return int(prediction[0])

# 创建 Gradio 接口
interface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(type="pil", label="Draw your digit here"),
    outputs="text",
    title="手写数字识别",
    description="在画布上直接绘制数字，模型将预测图像中的数字。"
)

interface.launch()