import numpy as np
import joblib
import gradio as gr
from PIL import Image
import os
from pinecone import Pinecone

# 创建 Pinecone 实例
pc = Pinecone(api_key="e97e2b32-22d1-4eb0-bbd3-02d95c2484ce")

# 连接到已创建的索引
index_name = "mnist-index"
index = pc.Index(index_name)

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

# 预处理图像
def preprocess_image(image):
    image = Image.fromarray(image)
    image = image.convert("L")  # 转换为灰度图
    image = image.resize((8, 8))  # 调整大小为 8x8
    image = np.array(image).flatten()  # 扁平化
    image = image / 16.0  # 归一化
    return image

# 定义使用 Pinecone 的预测函数
def predict_digit(image):
    # 预处理输入图像
    image = preprocess_image(image)
    vector = image.flatten().tolist()  # 将图像扁平化为列表

    # 使用 Pinecone 获取预测结果
    response = index.query(vector=vector, top_k=1)  # 查询 Pinecone 获取最近邻
    prediction = response['matches'][0]['id']  # 获取预测的数字 ID

    return int(prediction)

iface = gr.Interface(fn=predict_digit,
                     inputs=gr.Sketchpad(shape=(28, 28), label="在此处手写数字"),
                     outputs='label')

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