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

# 设置 Pinecone API 密钥
api_key = "c3e93ced-9abd-451a-afcd-73b1e5a9703e"

# 创建 Pinecone 实例
pc = Pinecone(api_key=api_key)

# 设置 Pinecone 索引
index_name = "mnist-index"
index = pc.Index(index_name)

def preprocess_image(image):
    if image is None:
        return None
    if not isinstance(image, np.ndarray):
        raise ValueError("image 不是一个 NumPy 数组")

    # 如果 image 是三维数组，将其转换为灰度图像
    if image.ndim == 3:
        image = np.mean(image, axis=2).astype(np.uint8)

    # 将图像转换为PIL Image对象，并指定模式为灰度（"L"）
    img = Image.fromarray(image, "L")
    
    # 调整图像大小为8x8
    img = img.resize((8, 8))
    
    # 转换为numpy数组并归一化到0到16之间
    img_array = (np.array(img) / 255.0) * 16
    return img_array.flatten()

def predict(image):
    processed_image = preprocess_image(image)
    if processed_image is None:
        return None
    
    # 使用 Pinecone 查询
    try:
        query_response = index.query(vector=processed_image.tolist(), top_k=11, include_metadata=True)
        neighbors = query_response['matches']
        
        if not neighbors:
            return "未找到邻居"
        
        # 获取最近的k个邻居的标签
        neighbor_labels = [match['metadata']['label'] for match in neighbors]
        
        # 使用多数投票法确定最终标签
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        
        return int(predicted_label)
    except Exception as e:
        return f"未知错误：{e}"

# 创建Gradio接口
iface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(label="Draw a digit here"),
    outputs=gr.Label(label="Prediction"),
    live=False,
    title="手写数字识别",
    description="请在下方的画板上绘制一个手写数字（0-9）"
)

# 启动接口
iface.launch(share=True, height=28, width=28)