Update app.py

app.py

@@ -3,11 +3,26 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-
 model_path = "DogClassifier2.1.keras"
 model = tf.keras.models.load_model(model_path)
 
-labels = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'American Spaniel', 'Basenji', 'Basset', 'Beagle', 
+# Define the core prediction function
+def predict_bmwX(image):
+    # Preprocess 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.expand_dims(image, axis=0)  # Add batch dimension
+
+    # Predict
+    prediction = model.predict(image)
+    
+    # Apply softmax to get probabilities for each class
+    prediction = tf.nn.softmax(prediction)
+    
+    # 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',
                    'Boston Terrier', 'Boxer', 'Bull Mastiff', 'Bull Terrier', 'Bulldog', 'Cairn', 'Chihuahua', 'Chinese Crested',
                    'Chow', 'Clumber','Cockapoo', 'Cocker', 'Collie', 'Corgi', 'Coyote', 'Dalmation', 'Dhole', 'Dingo', 'Doberman',
@@ -16,22 +31,19 @@ labels = ['Afghan', 'African Wild Dog', 'Airedale', 'American Hairless', 'Americ
                    'Lhasa', 'Malinois', 'Maltese', 'Maltese', 'Mex Hairless', 'Newfoundland', 'Pekinese', 'Pit Bull', 'Pomeranian',
                    'Poodle', 'Pug', 'Rhodesian', 'Rottweiler', 'Saint Bernard', 'Schnauzer', 'Scotch Terrier', 'Shar_Pei', 
                    'Shiba Inu', 'Shih-Tzu', 'Siberian Husky', 'Vizsla', 'Yorkie']
-
-def predict_regression(image):
-    # Preprocess image
-    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
-    image = image.resize((150, 150)).convert('L') #resize the image to 150x150 and converts it to gray scale
-    image = np.array(image)
-    print(image.shape)
-    # Predict
-    prediction = model.predict(image[None, ...])  # Assuming single regression value
-    confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
-    return confidences
+    
+    # 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))}
+    
+    # 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_regression,
+    fn=predict_bmwX,
     inputs=input_image, 
     outputs=gr.Label(),
     description="A simple MLP classification model for image classification using the MNIST dataset.")