Update app.py

app.py

@@ -11,10 +11,12 @@ def predict(image):
     # Convert image to a numpy array
     image = np.array(image)
     
-    # Resize the image to the expected shape (28, 28, 3)
+    # Resize the image to the expected shape (28, 28, 3) for RGB images
     image = tf.image.resize(image, (28, 28))  # Resize to 28x28 pixels
-    image = np.expand_dims(image, axis=-1)  # Add the channel dimension if grayscale
-    image = np.repeat(image, 3, axis=-1)  # Convert grayscale to RGB (if model was trained on RGB images)
+    
+    # Check if the image is grayscale (single channel), and convert to RGB if necessary
+    if image.shape[-1] == 1:  # If it's grayscale (single channel)
+        image = np.repeat(image, 3, axis=-1)  # Convert grayscale to RGB by repeating the channel
     
     # Normalize the image
     image = image / 255.0