import gradio as gr
import numpy as np
from collections import Counter
from pinecone_example import load_pinecone
import matplotlib.pyplot as plt
from PIL import Image

def predict_digit(stroke):
    image = np.array(stroke['composite'])
    plt.imshow(image, cmap='gray')
    plt.axis('off')
    plt.savefig('output_image.png')
    plt.close()

    image_array = image.flatten().reshape(-1, 64)[0].tolist()

    results = load_pinecone().query(vector=image_array, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']] if results['matches'] else []
    predicted_label = Counter(labels).most_common(1)[0][0] if labels else None

    return predicted_label

iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="绘制数字", crop_size=(8, 8), type="pil"),
    outputs="label",
    title="数字识别",
    description="绘制数字并识别其数字"
)

iface.launch()