Update app.py

app.py

@@ -19,22 +19,27 @@
 #     live=True
 # ).launch()
 
-def load_model_weights(model, checkpoint_path):
-    model.load_weights(checkpoint_path)
-
+import tensorflow as tf
 import gradio as gr
 import numpy as np
-from modelutil import create_model, load_model_weights
-
-checkpoint_path = './checkpoint'
 
-model = create_model()
-load_model_weights(model, checkpoint_path)
+# Load the CNN model from the .h5 file
+model = tf.keras.models.load_model('mnist_cnn_model.h5')
 
 def predict_digit(image):
-    predictions = model.predict(image.reshape(1, 28, 28))
-    return np.argmax(predictions)
+    # Preprocess the input image
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    image = image / 255.0  # Normalize pixel values
+
+    # Make predictions
+    predictions = model.predict(image)
 
+    # Get the predicted digit
+    predicted_digit = np.argmax(predictions)
+
+    return predicted_digit
+
+# Define Gradio interface
 gr.Interface(
     title="MNIST Digit Classifier by Papa Sega",
     fn=predict_digit, 
@@ -42,3 +47,4 @@ gr.Interface(
     outputs="number",
     live=True
 ).launch()
+