Update app.py

app.py

@@ -21,6 +21,9 @@ def predict_bmwX(image):
     # 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 shape: {prediction.shape}")
+    
     # Define class names
     class_names = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'American Spaniel', 'Basenji', 'Basset', 'Beagle', 
                    'Bearded Collie', 'Bermaise', 'Bichon Frise', 'Blenheim', 'Bloodhound', 'Bluetick', 'Border Collie', 'Borzoi',
@@ -32,6 +35,10 @@ def predict_bmwX(image):
                    'Poodle', 'Pug', 'Rhodesian', 'Rottweiler', 'Saint Bernard', 'Schnauzer', 'Scotch Terrier', 'Shar_Pei', 
                    '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)})."
+    
     # 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))}