Update app.py

app.py

@@ -8,34 +8,41 @@ from PIL import Image
 MODEL_PATH = "setosys_dogs_model.h5"
 model = tf.keras.models.load_model(MODEL_PATH)
 
-# Define image preprocessing function
-def preprocess_image(img):
+# Image preprocessing function
+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.expand_dims(img_array, axis=0)  # Add batch dimension
-    img_array = img_array / 255.0  # Normalize pixel values
+    img_array /= 255.0  # Normalize pixel values
     return img_array
 
-# Define the prediction function
-def predict_dog_breed(img):
+# Prediction function
+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)
-    class_idx = np.argmax(predictions)  # Get class index with highest probability
-    confidence = float(np.max(predictions))  # Get confidence score
-
-    # Define your class labels (replace with your actual class names)
-    class_labels = ["Labrador Retriever", "German Shepherd", "Golden Retriever", "Bulldog", "Poodle"]  
-    predicted_breed = class_labels[class_idx] if class_idx < len(class_labels) else "Unknown"
+    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"
+    
     return {predicted_breed: confidence}
 
-# Create a Gradio interface
+# Create Gradio interface
 interface = gr.Interface(
     fn=predict_dog_breed,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(),
     title="Dog Breed Classifier",
-    description="Upload an image of a dog to predict its breed.",
+    description="Upload an image of a dog to predict its breed."
 )
 
 # Launch the Gradio app