Update app.py

app.py

@@ -11,10 +11,10 @@ import gradio as gr
 import cv2
 
 def model(img):
-    model_dir = 'cataract'
-    class_names = ['Normal', 'Cataract']  # Cataract class labels
-    channel='RGB'
-    selected_model = 'cataract'
+    model_dir = 'pneumonia'
+    class_names = ['Normal eye', 'pneumonic lung']  # Pneumonia class labels
+    channel='L'
+    selected_model = 'pneumonia'
 
     architecture_path = os.path.join('models',model_dir, f'model_architecture_{selected_model}.pkl')
     weights_path = os.path.join('models',model_dir, f'model_weights_{selected_model}.pkl')
@@ -51,7 +51,7 @@ def model(img):
     predicted_class_index = 1 if (predictions[0]>0.5) else 0
     print(predicted_class_index)
 
-    predicted_class_label = class_names[predicted_class_index]
+    predicted_class_index = np.argmax(predictions[0])
 
     return predicted_class_label