Update app.py

app.py

@@ -7,14 +7,17 @@ from PIL import Image
 model = tf.keras.models.load_model("mnist_model.h5")
 
 def predict(image):
-    # Convert numpy array (from Sketchpad) to PIL
-    image = Image.fromarray(image.astype("uint8"), "RGB")
-    
-    # Convert to grayscale and resize
-    image = image.convert("L").resize((28,28))
+    # Extract numpy array if dict
+    if isinstance(image, dict) and "image" in image:
+        image = image["image"]
+
+    # Convert to PIL
+    image = Image.fromarray(image.astype("uint8")).convert("L").resize((28, 28))
+
+    # Normalize
     image = np.array(image) / 255.0
     image = image.reshape(1, 28, 28, 1)
-    
+
     # Predict
     prediction = model.predict(image)
     return str(np.argmax(prediction))