Update app.py

app.py

@@ -6,34 +6,23 @@ 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
     image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
     image = image.convert("RGB")  # Ensure the image is in RGB format
     image = image.resize((150, 150))  # Resize the image to 150x150
-    image = np.array(image) / 255.0  # Normalize the image
+    image = np.array(image)
     image = np.expand_dims(image, axis=0)  # Add batch dimension
 
     # Predict
-    raw_prediction = model.predict(image)
-    
-    # Debug statement to check the raw prediction values
-    print(f"Raw prediction: {raw_prediction}")
+    prediction = model.predict(image)
     
-    # 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]
+    # Apply softmax to get probabilities for each class
+    prediction = tf.nn.softmax(prediction)
     
-    # Debug statement to check the shape of the prediction
-    print(f"Prediction after softmax: {prediction}")
-    print(f"Sum of probabilities: {np.sum(prediction)}")
+    # Print prediction probabilities to the console
+    print("Prediction probabilities:", prediction)
     
     # Define class names
     class_names = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'American Spaniel', 'Basenji', 'Basset', 'Beagle', 
@@ -47,11 +36,21 @@ def predict_bmwX(image):
                    'Shiba Inu', 'Shih-Tzu', 'Siberian Husky', 'Vizsla', 'Yorkie']
     
     # Check if the number of predictions matches the number of class names
-    if len(prediction) != len(class_names):
-        return f"Error: Number of model outputs ({len(prediction)}) does not match number of class names ({len(class_names)})."
+    if len(prediction[0]) != len(class_names):
+        return f"Error: Number of model outputs ({len(prediction[0])}) does not match number of class names ({len(class_names)})."
     
+    # Apply threshold and set probabilities lower than 0.015 to 0.0
+    threshold = 0.015
+    prediction = np.array(prediction)
+    prediction[prediction < threshold] = 0.0
+    
+    # Recalculate the probabilities
+    total_probability = np.sum(prediction)
+    if total_probability > 0:
+        prediction = prediction / total_probability
+
     # Create a dictionary with the probabilities for each dog breed
-    prediction_dict = {class_names[i]: np.round(float(prediction[i]), 2) for i in range(len(class_names))}
+    prediction_dict = {class_names[i]: np.round(float(prediction[0][i]), 2) for i in range(len(class_names))}
     
     # Sort the dictionary by value in descending order and get the top 3 classes
     sorted_predictions = dict(sorted(prediction_dict.items(), key=lambda item: item[1], reverse=True))
@@ -63,5 +62,5 @@ iface = gr.Interface(
     fn=predict_bmwX,
     inputs=input_image, 
     outputs=gr.Label(),
-    description="A simple dog breed classification model.")
+    description="A simple MLP classification model for image classification using the MNIST dataset.")
 iface.launch(share=True)