Update app.py

app.py

@@ -6,6 +6,10 @@ from PIL import Image
 model_path = "DogClassifier2.2.keras"
 model = tf.keras.models.load_model(model_path)
 
+# Ensure the model's output layer is correct
+if not isinstance(model.layers[-1], tf.keras.layers.Softmax):
+    print("The last layer of the model is not a Softmax layer. The model might not be properly configured.")
+
 # Define the core prediction function
 def predict_bmwX(image):
     # Preprocess image
@@ -21,8 +25,11 @@ def predict_bmwX(image):
     # Debug statement to check the raw prediction values
     print(f"Raw prediction: {raw_prediction}")
     
-    # Apply softmax to get probabilities for each class
-    prediction = tf.nn.softmax(raw_prediction).numpy()[0]
+    # Apply softmax to get probabilities for each class if not already applied
+    if raw_prediction.shape[-1] != 1:
+        prediction = tf.nn.softmax(raw_prediction).numpy()[0]
+    else:
+        prediction = raw_prediction[0]
     
     # Debug statement to check the shape of the prediction
     print(f"Prediction after softmax: {prediction}")