Update app.py

app.py

@@ -10,23 +10,23 @@ model = tf.keras.models.load_model("denis_mnist_cnn_model.h5")
 def preprocess_image(image):
     # Convert PIL image to a tensor
     image = tf.convert_to_tensor(image)
-    
-    # Check if the image has a single channel (grayscale) and reshape if needed
-    if image.shape[-1] != 1:
-        # Convert to grayscale if the image is not in grayscale format (e.g., RGB)
-        image = tf.image.rgb_to_grayscale(image)
-    
+
     # Resize the image to 28x28 as expected by the model
     image = tf.image.resize(image, (28, 28))  # Resize to 28x28
 
-    # Convert grayscale to RGB (3 channels)
-    image = tf.image.grayscale_to_rgb(image)  # Convert grayscale to RGB (3 channels)
+    # If the image is RGB, convert it to grayscale
+    if image.shape[-1] == 3:
+        image = tf.image.rgb_to_grayscale(image)  # Convert RGB to grayscale
     
     # Normalize pixel values to [0, 1]
     image = image / 255.0
+    
+    # Convert grayscale to RGB (3 channels)
+    image = tf.image.grayscale_to_rgb(image)  # Convert grayscale to RGB (3 channels)
 
     # Add batch dimension (model expects batch of images)
     image = tf.expand_dims(image, axis=0)
+
     return image
 
 # Function to make predictions
@@ -34,7 +34,8 @@ def predict(image):
     image = preprocess_image(image)
     prediction = model.predict(image)  # Predict
     predicted_class = np.argmax(prediction, axis=-1)[0]  # Get the predicted class
-    return {"prediction": int(predicted_class)}
+    class_confidence = np.max(prediction, axis=-1)[0]  # Get the confidence score
+    return {"prediction": int(predicted_class), "confidence": class_confidence}
 
 # Create a Gradio interface
 interface = gr.Interface(fn=predict, inputs="image", outputs="json")