Update app.py

app.py

@@ -12,17 +12,21 @@ def predict_bmwX(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)
+    image = np.array(image) / 255.0  # Normalize the image
     image = np.expand_dims(image, axis=0)  # Add batch dimension
 
     # Predict
-    prediction = model.predict(image)
+    raw_prediction = model.predict(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(prediction)
+    prediction = tf.nn.softmax(raw_prediction).numpy()[0]
     
     # Debug statement to check the shape of the prediction
-    print(f"Prediction shape: {prediction.shape}")
+    print(f"Prediction after softmax: {prediction}")
+    print(f"Sum of probabilities: {np.sum(prediction)}")
     
     # Define class names
     class_names = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'American Spaniel', 'Basenji', 'Basset', 'Beagle', 
@@ -36,22 +40,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[0]) != len(class_names):
-        return f"Error: Number of model outputs ({len(prediction[0])}) does not match number of class names ({len(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)})."
     
     # Create a dictionary with the probabilities for each dog breed
-    prediction_dict = {class_names[i]: np.round(float(prediction[0][i]), 2) for i in range(len(class_names))}
+    prediction_dict = {class_names[i]: np.round(float(prediction[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))
         
     return sorted_predictions
 
-
 input_image = gr.Image()
 iface = gr.Interface(
     fn=predict_bmwX,
     inputs=input_image, 
     outputs=gr.Label(),
-    description="A simple MLP classification model for image classification using the MNIST dataset.")
+    description="A simple dog breed classification model.")
 iface.launch(share=True)