Spaces:
Sleeping
Sleeping
| import spaces | |
| import gradio as gr | |
| import numpy as np | |
| import tensorflow as tf | |
| from PIL import Image | |
| # Load the trained ANN model | |
| model = tf.keras.models.load_model("mnist_ann.keras") | |
| def preprocess(image): | |
| """Convert sketchpad/uploaded image to MNIST format: | |
| white digit on black, cropped to bounding box, centered in 28x28.""" | |
| if image is None: | |
| return None | |
| # Gradio Sketchpad returns a dict with 'composite' key | |
| if isinstance(image, dict): | |
| image = image.get("composite") | |
| if image is None: | |
| return None | |
| img = Image.fromarray(np.array(image).astype("uint8")).convert("L") | |
| arr = np.array(img).astype("float32") | |
| # MNIST = white digit on black background. | |
| # Sketchpad = black drawing on white, so invert if background is white. | |
| if arr.mean() > 127: | |
| arr = 255.0 - arr | |
| # --- Step 1: Crop to the bounding box of the drawing --- | |
| mask_rows = np.any(arr > 30, axis=1) | |
| mask_cols = np.any(arr > 30, axis=0) | |
| if not mask_rows.any() or not mask_cols.any(): | |
| return None # empty canvas | |
| r0, r1 = np.where(mask_rows)[0][[0, -1]] | |
| c0, c1 = np.where(mask_cols)[0][[0, -1]] | |
| arr = arr[r0:r1 + 1, c0:c1 + 1] | |
| # --- Step 2: Resize to fit in 20x20 (keeping aspect ratio) --- | |
| h, w = arr.shape | |
| scale = 20.0 / max(h, w) | |
| new_h = max(1, int(round(h * scale))) | |
| new_w = max(1, int(round(w * scale))) | |
| img_small = Image.fromarray(arr.astype("uint8")).resize( | |
| (new_w, new_h), Image.LANCZOS | |
| ) | |
| # --- Step 3: Paste in the center of a 28x28 black canvas --- | |
| canvas = np.zeros((28, 28), dtype="float32") | |
| top = (28 - new_h) // 2 | |
| left = (28 - new_w) // 2 | |
| canvas[top:top + new_h, left:left + new_w] = np.array(img_small) | |
| # Normalize to [0, 1] | |
| canvas = canvas / 255.0 | |
| return canvas.reshape(1, 28, 28) | |
| def predict_digit(image): | |
| arr = preprocess(image) | |
| if arr is None: | |
| return {"Draw a digit first!": 1.0} | |
| probs = model.predict(arr, verbose=0)[0] | |
| return {str(i): float(probs[i]) for i in range(10)} | |
| demo = gr.Interface( | |
| fn=predict_digit, | |
| inputs=gr.Sketchpad(label="Draw a digit (0-9)"), | |
| outputs=gr.Label(num_top_classes=3, label="Prediction"), | |
| title="MNIST Digit Classifier (ANN)", | |
| description="Assignment 3 - Model Deployment | Draw a handwritten digit and the ANN model will predict it.", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False) |