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

# Load the model directly
# reliable for Spaces where you upload the model file alongside the app
model = tf.keras.models.load_model("mnist_model.keras")

def predict_digit(image):
    if image is None:
        return None
    
    # 1. Convert to grayscale
    image = image.convert('L')
    
    # 2. Resize to 28x28 to match training data
    image = image.resize((28, 28))
    
    # 3. Invert colors (MNIST is white text on black background)
    # Most user uploads are black text on white background (paper), so we usually need to invert
    # We check mean pixel value; if > 127, it's likely a white background.
    if np.mean(image) > 127:
        image = ImageOps.invert(image)
        
    # 4. Convert to numpy array and normalize
    image_array = np.array(image) / 255.0
    
    # 5. Flatten to shape (1, 784) as expected by the Dense input layer
    image_array = image_array.reshape(1, 784)
    
    # Predict
    prediction = model.predict(image_array)
    
    # Return dictionary for Gradio Label output
    return {str(i): float(prediction[0][i]) for i in range(10)}

iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Label(num_top_classes=3, label="Predictions"),
    title="MNIST Digit Classifier",
    description="Upload an image of a handwritten digit (0-9) to see the prediction. Works best with a single digit centered in the image."
)

if __name__ == "__main__":
    iface.launch()