Update app.py

app.py

@@ -22,46 +22,19 @@ def predict_bmwX(image):
     # Apply softmax to get probabilities for each class
     prediction = tf.nn.softmax(prediction)
     
-    # Create a dictionary with the probabilities for each Pokemon
-    afghan = np.round(float(prediction[0][0]), 2)
-    africanWildDog = np.round(float(prediction[0][1]), 2)
-    beagle = np.round(float(prediction[0][2]), 2)
-    blenheim = np.round(float(prediction[0][3]), 2)
-    borderColie = np.round(float(prediction[0][4]), 2)
-    bostonTerrier = np.round(float(prediction[0][5]), 2)
-    chineseCrested = np.round(float(prediction[0][6]), 2)
-    cocker = np.round(float(prediction[0][7]), 2)
-    corgi = np.round(float(prediction[0][8]), 2)
-    dingo = np.round(float(prediction[0][9]), 2)
-    frenchBulldog = np.round(float(prediction[0][10]), 2)
-    germanShepard = np.round(float(prediction[0][11]), 2)
-    goldenRetriever = np.round(float(prediction[0][12]), 2)
-    pitBull = np.round(float(prediction[0][13]), 2)
-    rottweiler = np.round(float(prediction[0][14]), 2)
-    irishSpaniel = np.round(float(prediction[0][15]), 2)
-    labrador = np.round(float(prediction[0][16]), 2)
-    maltese = np.round(float(prediction[0][17]), 2)
-    newfoundland = np.round(float(prediction[0][18]), 2)
-    pomeranian = np.round(float(prediction[0][19]), 2)
-    poodle = np.round(float(prediction[0][20]), 2)
-    rhodesian = np.round(float(prediction[0][21]), 2)
-    saintBernard = np.round(float(prediction[0][22]), 2)
-    schnauzer = np.round(float(prediction[0][23]), 2)
-    scotchTerrier = np.round(float(prediction[0][24]), 2)
-    sharPei = np.round(float(prediction[0][25]), 2)
-    shibaInu = np.round(float(prediction[0][26]), 2)
-    siberianHusky = np.round(float(prediction[0][27]), 2)
-    yorkie = np.round(float(prediction[0][28]), 2)   
+    # Define class names
+    class_names = ['Afghan', 'African Wild Dog', 'Beagle', 'Blenheim', 'Border Collie', 'Boston Terrier', 'Chinese Crested', 
+                   'Cocker', 'Corgi', 'Dingo', 'French Bulldog', 'German Shepard', 'Golden Retriever', 'Pit Bull', 
+                   'Rottweiler', 'Irish Spaniel', 'Labrador', 'Maltese', 'Newfoundland', 'Pomeranian', 'Poodle', 
+                   'Rhodesian', 'Saint Bernard', 'Schnauzer', 'Scotch Terrier', 'Shar Pei', 'Shiba Inu', 'Siberian Husky', 'Yorkie']
     
+    # 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
+    top_3 = dict(sorted(prediction_dict.items(), key=lambda item: item[1], reverse=True)[:3])
         
-    return {'Afghan': afghan, 'African Wild Dog': africanWildDog, 'Beagle': beagle, 'Blenheim': blenheim,
-            'Border Collie': borderColie, 'Boston Terrier': bostonTerrier, 'Chinese Crested': chineseCrested, 'Cocker': cocker,
-           'Corgi': corgi, 'Dingo': dingo, 'French Bulldog': frenchBulldog, 'German Shepard': germanShepard,
-           'GoldenRetriever': goldenRetriever, 'Pit Bull': pitBull, 'Rottweiler': rottweiler, 'Irish Spaniel': irishSpaniel,
-           'Labrador': labrador, 'Maltese': maltese, 'Newfoundland': newfoundland, 'Pomeranian': pomeranian,
-           'Poodle': poodle, 'Rhodesian': rhodesian, 'Saint Bernard': saintBernard, 'Schnauzer': schnauzer,
-           'Scotch Terrier': scotchTerrier, 'Shar Pei': sharPei, 'Shiba Inu': shibaInu, 'Siberian Husky': siberianHusky,
-           'Yorkie': yorkie}
+    return top_3
 
 
 input_image = gr.Image()
@@ -69,5 +42,5 @@ iface = gr.Interface(
     fn=predict_bmwX,
     inputs=input_image, 
     outputs=gr.Label(),
-    description="A simple mlp classification model for image classification using the mnist dataset.")
-iface.launch(share=True)
+    description="A simple MLP classification model for image classification using the MNIST dataset.")
+iface.launch(share=True)