Update app.py

app.py

@@ -2,19 +2,24 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 from tensorflow.keras.preprocessing import image
+from tensorflow.keras.applications.efficientnet_v2 import EfficientNetV2S, preprocess_input
 from PIL import Image
 
 # Load the trained model
 MODEL_PATH = "setosys_dogs_model.h5"
 model = tf.keras.models.load_model(MODEL_PATH)
 
-# Image preprocessing function
+# 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
+
+# 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 = image.img_to_array(img)
+    img_array = np.array(img)
+    img_array = EfficientNetV2S.preprocess_input(img_array)  # EfficientNetV2 preprocessing
     img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
-    img_array /= 255.0  # Normalize pixel values
     return img_array
 
 # Prediction function
@@ -24,15 +29,9 @@ def predict_dog_breed(img: Image.Image) -> dict:
     predictions = model.predict(img_array)
     class_idx = np.argmax(predictions)  # Index of the highest prediction probability
     confidence = float(np.max(predictions))  # Confidence score
-
-    # Get class labels from the model
-    class_labels = model.classes_ if hasattr(model, 'classes_') else None
     
-    # If class labels are available, return the predicted breed with confidence score
-    if class_labels is not None:
-        predicted_breed = class_labels[class_idx]
-    else:
-        predicted_breed = "Unknown"
+    # Get predicted breed and its confidence score
+    predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown"
     
     return {predicted_breed: confidence}