Update app.py

app.py

@@ -26,8 +26,8 @@ def preprocess_image(image, target_size=(224, 224)):
     return img_array
 
 # ===== Disease Label Mappings =====
-# The five disease labels you are focusing on.
-disease_labels = {
+# Model A labels (for example)
+disease_labels_a = {
     0: "Tomato Bacterial Spot",
     1: "Tomato Early Blight",
     2: "Tomato Late Blight",
@@ -35,18 +35,27 @@ disease_labels = {
     4: "Tomato Yellow Leaf Curl Virus"
 }
 
+# Model B labels (based on your training dataset)
+disease_labels_b = {
+    0: "Tomato___Target_Spot",
+    1: "Tomato___Bacterial_spot",
+    2: "Tomato___Early_blight",
+    3: "Tomato___healthy",
+    4: "Tomato___Late_blight"
+}
+
 # ===== Prediction Functions =====
 def predict_model_a(image):
     img = preprocess_image(image)
     pred = model_a.predict(img)
     predicted_class = np.argmax(pred)
-    return disease_labels.get(predicted_class, "Unknown result")
+    return disease_labels_a.get(predicted_class, "Unknown result")
 
 def predict_model_b(image):
     img = preprocess_image(image)
     pred = model_b.predict(img)
     predicted_class = np.argmax(pred)
-    return disease_labels.get(predicted_class, "Unknown result")
+    return disease_labels_b.get(predicted_class, "Unknown result")
 
 def predict_classifier(image):
     img = preprocess_image(image)
@@ -87,7 +96,7 @@ def call_openassistant(prompt):
 
 # ===== AI Assistant Functions =====
 def ai_assistant_v1(image, prediction):
-    # Use Llama 2-7B Chat (Model A versions)
+    # Use Llama 2-7B Chat (for Model A versions)
     if "Tomato" not in prediction:
         prompt = (
             "You are an agricultural advisor. The tomato crop appears healthy. "
@@ -101,7 +110,7 @@ def ai_assistant_v1(image, prediction):
     return call_llama2(prompt)
 
 def ai_assistant_v2(image, prediction):
-    # Use OpenAssistant (Model B versions)
+    # Use OpenAssistant (for Model B versions)
     if "Tomato" not in prediction:
         prompt = (
             "You are an agricultural advisor. The tomato crop appears healthy. "