import gradio as gr  # 用于创建图形用户界面（GUI）
import pickle  # 用于加载保存的模型
import cv2  # 用于图像处理

import matplotlib.pyplot as plt

# ######## 1. 加载模型 ########
# 打开并加载保存的 KNN 模型，使用 'rb' 模式读取文件
with open('best_knn_model.pkl', 'rb') as file:
    knn = pickle.load(file)

# ######## 2. 定义预测函数 ########
# 处理输入的图像并进行预测
def predict(image):
    if image is None:  # 检查图像是否为空
        return None  # 如果为空，返回 None
    # 转换图像为灰度，保持模型输入的一致性
    image = cv2.cvtColor(image["composite"], cv2.COLOR_RGBA2GRAY)
    plt.imsave('image1.png', image, cmap='gray') # 把获取的图像保存到本地，以便检查获取的图像是否正确
    
    # 调整图像大小为 8x8 并展平成一维数组，符合模型输入要求
    image = cv2.resize(image, (8, 8), interpolation=cv2.INTER_AREA)
    plt.imsave('image2.png', image, cmap='gray') # 把调整大小后的图像保存到本地，以便检查调整大小后的图像是否正确
    
    image = image.ravel()
    # 使用加载的 KNN 模型进行预测，返回结果并转化为字符串
    return str(knn.predict([image])[0])

# ######## 3. 构建 Gradio 界面 ########
# 创建一个界面，使用图像编辑器作为输入，文本标签作为输出
gr.Interface(fn=predict, inputs=gr.ImageEditor(), outputs="label").launch()