import gradio as gr
import numpy as np
from PIL import Image
from tensorflow import keras

model = keras.models.load_model("Digit.keras")

def preprocess_image(img):
    if img is None:
        raise ValueError("No image provided.")

    # If Sketchpad returns a dict, try common keys
    if isinstance(img, dict):
        if "image" in img and img["image"] is not None:
            img = img["image"]
        elif "composite" in img and img["composite"] is not None:
            img = img["composite"]
        else:
            raise ValueError("No image data found in dictionary.")

    # If still a NumPy array, convert to PIL
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype("uint8"))

    # --- MNIST preprocessing ---
    img = img.convert("L").resize((28, 28))
    arr = np.array(img).astype("float32")
    if arr.mean() > 127:  # invert if white background
        arr = 255 - arr
    arr = arr / 255.0
    return arr[np.newaxis, ..., np.newaxis]  # (1,28,28,1)

def predict(img):
    try:
        x = preprocess_image(img)
        probs = model.predict(x, verbose=0)[0]
        pred = int(np.argmax(probs))
        return str(pred), {str(i): float(probs[i]) for i in range(10)}
    except Exception as e:
        return f"Error: {e}", {}


with gr.Blocks(title="MNIST Digit Classifier") as demo:
    gr.Markdown("## ✍️ Draw a digit (0–9)")
    with gr.Row():
        with gr.Column():
            canvas = gr.Sketchpad(
                canvas_size=(280, 280),
                type="pil",                 # <-- forces PIL, avoids dicts
                label="Draw a digit here"
            )
            btn = gr.Button("Predict")
        with gr.Column():
            pred_txt = gr.Textbox(label="Predicted Digit", interactive=False)
            probs = gr.Label(label="Class Probabilities", num_top_classes=10)

    btn.click(predict, inputs=canvas, outputs=[pred_txt, probs])

demo.launch(share=True)