Update app.py

app.py

@@ -25,22 +25,21 @@ def predict(image):
     custom_cnn_pred = custom_cnn_model.predict(np.expand_dims(processed_image, axis=0))[0]
     
     # Get the predicted labels
-    vgg_label = np.argmax(vgg_pred,axis=1)
-    mobilenet_label = np.argmax(mobilenet_pred,axis=1)
-    custom_cnn_label = np.argmax(custom_cnn_pred,axis=1)
+    vgg_label = np.argmax(vgg_pred)
+    mobilenet_label = np.argmax(mobilenet_pred)
+    custom_cnn_label = np.argmax(custom_cnn_pred)
     
     label_map = {0: 'Black Sea Sprat', 1: 'Gilt-Head Bream', 2: 'Hourse Mackerel', 
              3: 'Red Mullet', 4: 'Red Sea Bream', 5: 'Sea Bass', 
              6: 'Shrimp', 7: 'Striped Red Mullet', 8: 'Trout'}
-    label1 = [label_map[label] for label in custom_cnn_label][0]
-    label2 = [label_map[label] for label in vgg_label][0]
-    label3 = [label_map[label] for label in mobilenet_label][0]
-
-    
     
+    label1 = label_map[custom_cnn_label]
+    label2 = label_map[vgg_label]
+    label3 = label_map[mobilenet_label]
     
     return label1, label2, label3
 
+
 # Create the Gradio interface
 inputs = gr.components.Image()
 outputs=[