draw_digits / app.py
JasonKishore's picture
Update app.py
fd5ef98 verified
Raw
History Blame Contribute Delete
1.95 kB
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)