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

MODEL_PATH = "mnist_cnn.h5"

model = tf.keras.models.load_model(MODEL_PATH)


def preprocess(image):
    if image is None:
        return None

    # Ensure PIL
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image.astype("uint8"))

    # Convert to grayscale
    image = image.convert("L")

    # Resize
    image = image.resize((28, 28))

    img = np.array(image).astype("float32")

    # Invert colors (MNIST style)
    img = 255 - img

    # Remove background noise
    img[img < 40] = 0

    # Normalize
    img /= 255.0

    # Add dims
    img = np.expand_dims(img, axis=-1)
    img = np.expand_dims(img, axis=0)

    return img


def predict(image):
    img = preprocess(image)
    preds = model.predict(img, verbose=0)[0]
    return {str(i): float(preds[i]) for i in range(10)}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(
        label="Upload a digit image (white digit on dark background)"
    ),
    outputs=gr.Label(num_top_classes=3),
    title="MNIST Handwritten Digit Classifier",
    description=(
        "Upload an image of a handwritten digit.\n\n"
        "Tips:\n"
        "- Dark background\n"
        "- Light digit\n"
        "- Centered and large"
    ),
)

demo.launch()