Upload app.py

app.py

@@ -20,18 +20,27 @@ model = tf.keras.models.load_model("number_recognition_model_colab.keras")
 
 
 def predict(img):
-    # Convert the input image to a NumPy array if needed
-    if not isinstance(img, np.ndarray):
-        img = np.array(img)
-
-    # Print the type and shape of the image
-    print(type(img), img.shape)
-
-    # Resize the image
-    img = cv2.resize(img, (img_size, img_size))
-    img = img.reshape(1, img_size, img_size, 1)
-    preds = model.predict(img)[0]
-    return {label: float(pred) for label, pred in zip(labels, preds)}
+    try:
+        # Convert the input image to a NumPy array if needed
+        if not isinstance(img, np.ndarray):
+            img = np.array(img)
+
+        # Ensure the image has a single channel (grayscale)
+        if img.ndim == 2:
+            img = np.expand_dims(img, axis=-1)
+
+        # Print the type and shape of the image
+        print(type(img), img.shape)
+
+        # Resize the image
+        img = cv2.resize(img, (img_size, img_size))
+        img = img.reshape(1, img_size, img_size, 1)
+        preds = model.predict(img)[0]
+        return {label: float(pred) for label, pred in zip(labels, preds)}
+    except Exception as e:
+        # Print the exception to the console
+        print(f"Error during prediction: {e}")
+        return {"Error": str(e)}
 
 
 label = gr.Label(num_top_classes=3)