Update app.py

app.py

@@ -25,11 +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)
-    mobilenet_label = np.argmax(mobilenet_pred)
-    custom_cnn_label = np.argmax(custom_cnn_pred)
+    vgg_label = np.argmax(vgg_pred,axis=1)
+    mobilenet_label = np.argmax(mobilenet_pred,axis=1)[0]
+    custom_cnn_label = np.argmax(custom_cnn_pred,axis=1)[0]
     
-    return custom_cnn_label, vgg_label, mobilenet_label
+    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]
+
+    
+    
+    
+    return label1, label2, label3
 
 # Create the Gradio interface
 inputs = gr.components.Image()