Update app.py

app.py

@@ -1,7 +1,6 @@
 import numpy as np
-import pandas as pd
 import gradio as gr
-from matplotlib import pyplot as plt
+from PIL import Image
 
 def ReLU(Z):
     return np.maximum(Z, 0)
@@ -32,16 +31,16 @@ def make_predictions(X, W1, b1, W2, b2):
     predictions = get_predictions(A2)
     return predictions
 
-def predict_digit(input_data):
+def predict_digit(img):
     # Load the trained parameters
     params = np.load("trained_params.npz", allow_pickle=True)
     W1, b1, W2, b2 = params["W1"], params["b1"], params["W2"], params["b2"]
 
-    # Convert the uploaded image to grayscale and resize it to (28, 28)
-    img = Image.fromarray(input_data).convert("L").resize((28, 28))
+    # Convert the sketchpad drawing to grayscale and resize it to (28, 28)
+    img_pil = Image.fromarray(np.uint8(img * 255)).convert("L").resize((28, 28))
 
     # Convert the image to a NumPy array and normalize it
-    X = np.array(img).reshape((784, 1)) / 255.
+    X = np.array(img_pil).reshape((784, 1)) / 255.
 
     # Get the prediction
     prediction = make_predictions(X, W1, b1, W2, b2)
@@ -50,7 +49,7 @@ def predict_digit(input_data):
 
 iface = gr.Interface(
     fn=predict_digit,
-    inputs=gr.inputs.Image(shape=(28, 28)),
+    inputs="sketchpad",
     outputs=gr.outputs.Label(num_top_classes=3),
     live=True,
     capture_session=True,