Update app.py

app.py

@@ -42,7 +42,17 @@ class_names = ['BacterialBlights', 'Healthy', 'Mosaic', 'RedRot', 'Rust', 'Yello
 # Load AI response generator (using a local GPT pipeline or OpenAI's GPT-3/4 API)
 ai_pipeline = pipeline("text-generation", model="gpt2", tokenizer="gpt2")
 
-# Function to predict disease type from an image
+# Knowledge base for sugarcane diseases (example data from the website)
+knowledge_base = {
+    'BacterialBlights': "Bacterial blights cause water-soaked lesions on leaves, leading to yellowing and withering. To manage, apply copper-based fungicides and ensure proper drainage.",
+    'Mosaic': "Mosaic disease results in streaked and mottled leaves, reducing photosynthesis. Use disease-resistant varieties and control aphids to prevent spread.",
+    'RedRot': "Red rot is identified by reddening and rotting of stalks. Remove infected plants and treat soil with appropriate fungicides.",
+    'Rust': "Rust appears as orange pustules on leaves. Apply systemic fungicides and maintain optimal field conditions to reduce spread.",
+    'Yellow': "Yellowing indicates nutrient deficiencies or initial disease stages. Test soil and provide balanced fertilizers.",
+    'Healthy': "The sugarcane crop is healthy. Continue regular monitoring and good agronomic practices."
+}
+
+# Update the predict_disease function
 def predict_disease(image):
     # Apply transformations to the image
     img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
@@ -55,13 +65,13 @@ def predict_disease(image):
     # Get the predicted label
     predicted_label = class_names[predicted_class.item()]
 
-    # Generate a detailed response for the detected disease
-    prompt = f"The detected sugarcane disease is '{predicted_label}'. Provide detailed advice for managing this condition."
-    ai_response = ai_pipeline(prompt, max_length=100, num_return_sequences=1, truncation=True)[0]['generated_text']
-    
-    # Post-process the AI response to ensure it ends with a complete sentence
-    if not ai_response.endswith(('.', '!', '?')):
-        ai_response = ai_response.rsplit('.', 1)[0] + '.'
+    # Retrieve response from knowledge base
+    if predicted_label in knowledge_base:
+        detailed_response = knowledge_base[predicted_label]
+    else:
+        # Fallback to AI-generated response
+        prompt = f"The detected sugarcane disease is '{predicted_label}'. Provide detailed advice for managing this condition."
+        detailed_response = ai_pipeline(prompt, max_length=100, num_return_sequences=1, truncation=True)[0]['generated_text']
 
     # Create a styled HTML output
     output_message = f"""
@@ -73,13 +83,13 @@ def predict_disease(image):
     if predicted_label != "Healthy":
         output_message += f"""
         <p style='font-size: 16px; color: #757575;'>
-            {ai_response}
+            {detailed_response}
         </p>
         """
     else:
         output_message += f"""
         <p style='font-size: 16px; color: #757575;'>
-            The sugarcane crop is <strong>healthy</strong>. Keep monitoring for potential risks.
+            {detailed_response}
         </p>
         """