Update app.py

app.py

@@ -9,16 +9,15 @@ from PIL import Image
 MODEL_PATH = "setosys_dogs_model.h5"
 model = tf.keras.models.load_model(MODEL_PATH)
 
-# Ideally, you would have access to train_generator's class_indices, or you can load them manually
-# Example if you manually define class labels
-class_labels = ["Labrador Retriever", "German Shepherd", "Golden Retriever", "Bulldog", "Poodle"]  # Adjust with actual labels
+# Get class labels from the model (assuming the model has a 'class_indices' attribute)
+class_labels = list(model.class_indices.keys())  # Fetch class labels from the model
 
 # Image preprocessing function using EfficientNetV2S
 def preprocess_image(img: Image.Image) -> np.ndarray:
     """Preprocess the image to match the model's input requirements."""
     img = img.resize((224, 224))  # Resize image to model input size
     img_array = np.array(img)
-    img_array = preprocess_input(img_array)  # EfficientNetV2 preprocessing (correct usage)
+    img_array = preprocess_input(img_array)  # EfficientNetV2 preprocessing
     img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
     return img_array
 
@@ -27,6 +26,10 @@ def predict_dog_breed(img: Image.Image) -> dict:
     """Predict the breed of the dog in the uploaded image."""
     img_array = preprocess_image(img)
     predictions = model.predict(img_array)
+    
+    # Check the shape of the predictions to make sure the output is correct
+    print("Predictions Shape:", predictions.shape)
+    
     class_idx = np.argmax(predictions)  # Index of the highest prediction probability
     confidence = float(np.max(predictions))  # Confidence score