Update app.py

app.py

@@ -1,61 +1,123 @@
 import gradio as gr
 import pandas as pd
 import os
-import random
-from PIL import Image
 import torch
-from torchvision import transforms
+from torchvision import models, transforms
+from PIL import Image
+import numpy as np
+import json
+
+# --- Model Setup ---
+def load_model():
+    # Load a pre-trained ResNet model
+    model = models.resnet50(pretrained=False)
+    
+    # Modify the final layer for our number of classes
+    num_classes = 38  # PlantVillage has 38 classes
+    model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
+    
+    # Load the trained weights
+    try:
+        model.load_state_dict(torch.load('plant_disease_model.pth', map_location=torch.device('cpu')))
+        print("Model loaded successfully")
+    except:
+        print("Model weights not found, using untrained model for demonstration")
+    
+    model.eval()
+    return model
+
+# --- Data Loading ---
+def load_class_names():
+    try:
+        with open('class_names.json', 'r') as f:
+            class_names = json.load(f)
+        print("Class names loaded successfully")
+        return class_names
+    except:
+        print("Class names file not found, using default classes")
+        # Default class names from PlantVillage dataset
+        return [
+            "Apple___Apple_scab", "Apple___Black_rot", "Apple___Cedar_apple_rust", "Apple___healthy",
+            "Cherry___healthy", "Cherry___Powdery_mildew", "Corn___Cercospora_leaf_spot", 
+            "Corn___Common_rust", "Corn___healthy", "Corn___Northern_Leaf_Blight",
+            "Grape___Black_rot", "Grape___Esca_(Black_Measles)", "Grape___healthy", 
+            "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)", "Orange___Haunglongbing_(Citrus_greening)",
+            "Peach___Bacterial_spot", "Peach___healthy", "Pepper,_bell___Bacterial_spot", 
+            "Pepper,_bell___healthy", "Potato___Early_blight", "Potato___healthy", 
+            "Potato___Late_blight", "Squash___Powdery_mildew", "Strawberry___healthy",
+            "Strawberry___Leaf_scorch", "Tomato___Bacterial_spot", "Tomato___Early_blight", 
+            "Tomato___healthy", "Tomato___Late_blight", "Tomato___Leaf_Mold", 
+            "Tomato___Septoria_leaf_spot", "Tomato___Spider_mites Two-spotted_spider_mite", 
+            "Tomato___Target_Spot", "Tomato___Tomato_mosaic_virus", 
+            "Tomato___Tomato_Yellow_Leaf_Curl_Virus"
+        ]
 
-# --- Load CSV Data ---
 def load_treatments_data():
     try:
-        context = pd.read_csv('./crop_diseases_treatments.csv')
+        context = pd.read_csv('crop_diseases_treatments.csv')
         print("CSV file loaded successfully.")
         return context
     except FileNotFoundError:
         print("Error: crop_diseases_treatments.csv not found.")
         # Create a minimal dataframe for demonstration
         return pd.DataFrame({
-            'Crop': ['Tomato', 'Apple', 'Corn', 'Rice', 'Potato'],
-            'Disease': ['Early Blight', 'Apple Scab', 'Common Rust', 'Bacterial Leaf Blight', 'Late Blight'],
-            'Symptoms': ['Yellow spots on leaves', 'Dark scab-like lesions', 'Rust-colored pustules', 'Yellow-orange lesions', 'Dark water-soaked spots'],
-            'Treatment': ['Remove affected leaves', 'Prune affected branches', 'Remove infected plants', 'Drain fields', 'Apply fungicide early'],
-            'Medicine/Chemical Control': ['Copper fungicide', 'Sulfur spray', 'Propiconazole', 'Streptomycin sulfate', 'Chlorothalonil']
+            'Crop': ['Tomato', 'Apple', 'Corn', 'Potato', 'Grape', 'Cherry', 'Peach', 'Strawberry'],
+            'Disease': ['Early Blight', 'Apple Scab', 'Common Rust', 'Late Blight', 'Black Rot', 'Powdery Mildew', 'Bacterial Spot', 'Leaf Scorch'],
+            'Symptoms': [
+                'Brown spots with concentric rings on leaves', 
+                'Olive-green to brown spots on leaves and fruit',
+                'Rust-colored pustules on leaves',
+                'Dark water-soaked spots on leaves that turn brown',
+                'Reddish-brown spots on leaves and fruit',
+                'White powdery coating on leaves and stems',
+                'Small brown spots on leaves and fruit',
+                'Scorched appearance on leaf margins'
+            ],
+            'Treatment': [
+                'Remove affected leaves, improve air circulation, rotate crops',
+                'Prune affected branches, remove fallen leaves, apply fungicide',
+                'Remove infected plants, apply fungicide early in season',
+                'Remove infected plants, avoid overhead watering, apply fungicide',
+                'Prune infected areas, remove mummified fruit, apply fungicide',
+                'Improve air circulation, apply fungicide, remove infected parts',
+                'Copper-based sprays, crop rotation, remove infected plants',
+                'Ensure proper watering, add mulch, improve soil drainage'
+            ],
+            'Medicine/Chemical Control': [
+                'Chlorothalonil, Mancozeb, Copper fungicides',
+                'Captan, Myclobutanil, Sulfur sprays',
+                'Propiconazole, Azoxystrobin, Mancozeb',
+                'Chlorothalonil, Mancozeb, Copper-based fungicides',
+                'Captan, Myclobutanil, Mancozeb',
+                'Sulfur, Potassium bicarbonate, Neem oil',
+                'Copper hydroxide, Streptomycin sulfate',
+                'Calcium nitrate sprays, Fungicides with Captan'
+            ]
         })
 
-# --- Disease Classification ---
-def prepare_disease_classifier():
-    # Set up image transformation
-    transform = transforms.Compose([
+# --- Image Processing ---
+def prepare_image_transform():
+    return transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
-    
-    # Example class names (to be replaced with actual classes from the dataset)
-    class_names = [
-        "Apple___Apple_scab", 
-        "Apple___Black_rot", 
-        "Apple___Cedar_apple_rust", 
-        "Apple___healthy", 
-        "Corn_(maize)___Cercospora_leaf_spot", 
-        "Corn_(maize)___Common_rust", 
-        "Corn_(maize)___healthy", 
-        "Tomato___Early_blight", 
-        "Tomato___Late_blight", 
-        "Tomato___healthy"
-    ]
-    
-    return transform, class_names
 
-# Function to classify plant disease from image
-def classify_disease(image):
-    transform, class_names = prepare_disease_classifier()
+# --- Disease Classification ---
+def classify_disease(image, model, class_names, transform):
+    if image is None:
+        return None, None
     
     # Process the image
-    if image is not None:
-        # For demonstration, we'll randomly select a class
-        class_idx = random.randint(0, len(class_names) - 1)
+    img_tensor = transform(image).unsqueeze(0)
+    
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        _, predicted = torch.max(outputs, 1)
+        class_idx = predicted.item()
+    
+    if class_idx < len(class_names):
         class_name = class_names[class_idx]
         
         # Extract crop and disease from class name
@@ -65,9 +127,9 @@ def classify_disease(image):
         
         return crop, disease
     
-    return None, None
+    return "Unknown", "Unknown"
 
-# Function to find treatment in the database
+# --- Treatment Lookup ---
 def find_treatment(crop, disease, df):
     if disease.lower() == "healthy":
         return {
@@ -108,11 +170,11 @@ def find_treatment(crop, disease, df):
             "crop": crop,
             "disease": disease,
             "symptoms": "Unknown",
-            "treatment": "No specific treatment information found",
-            "medicine": "Consult a plant pathologist"
+            "treatment": "No specific treatment information found. General advice: Remove affected parts, ensure proper spacing for air circulation, and consider organic or chemical fungicides appropriate for your region.",
+            "medicine": "Consult a local agricultural extension office or plant pathologist for specific recommendations for your region."
         }
 
-# Simple function to answer questions based on the database
+# --- Q&A Function ---
 def answer_question(question, df):
     question = question.lower()
     
@@ -136,35 +198,40 @@ def answer_question(question, df):
         if not matches.empty:
             match = matches.iloc[0]
             return f"""
-            Crop: {match['Crop']}
-            Disease: {match['Disease']}
-            Symptoms: {match['Symptoms']}
-            Treatment: {match['Treatment']}
-            Medicine: {match['Medicine/Chemical Control']}
+            ## {match['Crop']} - {match['Disease']}
+            
+            **Symptoms:**
+            {match['Symptoms']}
+            
+            **Treatment:**
+            {match['Treatment']}
+            
+            **Recommended Products:**
+            {match['Medicine/Chemical Control']}
             """
     
     # If we found only crop
     elif crop_match:
         matches = df[df['Crop'] == crop_match]
         if not matches.empty:
-            response = f"Information about diseases affecting {crop_match}:\n\n"
+            response = f"## Common Diseases Affecting {crop_match}\n\n"
             for _, row in matches.iterrows():
-                response += f"Disease: {row['Disease']}\n"
-                response += f"Symptoms: {row['Symptoms']}\n"
-                response += f"Treatment: {row['Treatment']}\n"
-                response += f"Medicine: {row['Medicine/Chemical Control']}\n\n"
+                response += f"### {row['Disease']}\n"
+                response += f"**Symptoms:** {row['Symptoms']}\n\n"
+                response += f"**Treatment:** {row['Treatment']}\n\n"
+                response += f"**Products:** {row['Medicine/Chemical Control']}\n\n"
             return response
     
     # If we found only disease
     elif disease_match:
         matches = df[df['Disease'] == disease_match]
         if not matches.empty:
-            response = f"Information about {disease_match} affecting various crops:\n\n"
+            response = f"## {disease_match} in Different Crops\n\n"
             for _, row in matches.iterrows():
-                response += f"Crop: {row['Crop']}\n"
-                response += f"Symptoms: {row['Symptoms']}\n"
-                response += f"Treatment: {row['Treatment']}\n"
-                response += f"Medicine: {row['Medicine/Chemical Control']}\n\n"
+                response += f"### {row['Crop']}\n"
+                response += f"**Symptoms:** {row['Symptoms']}\n\n"
+                response += f"**Treatment:** {row['Treatment']}\n\n"
+                response += f"**Products:** {row['Medicine/Chemical Control']}\n\n"
             return response
     
     # General search
@@ -195,80 +262,107 @@ def answer_question(question, df):
         relevant_rows.sort(key=lambda x: x[0], reverse=True)
         top_matches = relevant_rows[:3]  # Get top 3 matches
         
-        response = "Here's what I found based on your question:\n\n"
+        response = "## Here's what I found based on your question:\n\n"
         for _, row in top_matches:
-            response += f"Crop: {row['Crop']}\n"
-            response += f"Disease: {row['Disease']}\n"
-            response += f"Symptoms: {row['Symptoms']}\n"
-            response += f"Treatment: {row['Treatment']}\n"
-            response += f"Medicine: {row['Medicine/Chemical Control']}\n\n"
+            response += f"### {row['Crop']} - {row['Disease']}\n"
+            response += f"**Symptoms:** {row['Symptoms']}\n\n"
+            response += f"**Treatment:** {row['Treatment']}\n\n"
+            response += f"**Products:** {row['Medicine/Chemical Control']}\n\n"
         return response
     
-    return "I couldn't find specific information related to your question in the database. Please try asking about specific crops or diseases."
+    return "I couldn't find specific information related to your question in my database. Please try asking about specific crops or diseases, or upload an image for analysis."
 
-# Function to process image uploads
-def process_image(image, df):
+# --- Image Processing Function ---
+def process_image(image, model, class_names, transform, df):
     if image is None:
         return "Please upload an image to analyze.", None, None, None, None
     
     # Identify the crop and disease
-    crop, disease = classify_disease(image)
+    crop, disease = classify_disease(image, model, class_names, transform)
+    
+    if crop is None or disease is None:
+        return "Unable to analyze the image. Please try a clearer image of a plant leaf.", None, None, None, None
     
     # Get treatment information
     treatment_info = find_treatment(crop, disease, df)
     
     # Create analysis text
     if disease.lower() == "healthy":
-        analysis = f"Good news! This {crop} plant appears to be healthy."
+        analysis = f"## Good news! \nThis {crop} plant appears to be healthy."
     else:
-        analysis = f"Analysis: This {crop} plant appears to be affected by {disease}.\n\n"
-        analysis += f"Typical symptoms of {disease} include {treatment_info['symptoms']}.\n\n"
-        analysis += f"Recommended treatment: {treatment_info['treatment']}"
+        analysis = f"## Analysis Results\n\nThis {crop} plant appears to be affected by **{disease}**.\n\n"
+        analysis += f"**Typical symptoms:**\n{treatment_info['symptoms']}\n\n"
+        analysis += f"**Recommended treatment:**\n{treatment_info['treatment']}\n\n"
+        analysis += f"**Recommended products:**\n{treatment_info['medicine']}"
     
     return analysis, treatment_info["crop"], treatment_info["disease"], treatment_info["treatment"], treatment_info["medicine"]
 
-# Main function to set up the Gradio interface
+# --- Main Function ---
 def main():
-    # Load data
+    # Load model and data
+    model = load_model()
+    class_names = load_class_names()
+    transform = prepare_image_transform()
     df = load_treatments_data()
     
-    # Create a simplified Gradio interface
-    with gr.Blocks(title="Plant Disease Assistant") as app:
-        gr.Markdown("# Plant Disease Treatment Assistant")
-        gr.Markdown("Upload a plant image or ask a question to get disease and treatment information")
+    # Create Gradio interface
+    with gr.Blocks(title="Plant Disease Assistant", css="footer {visibility: hidden}") as app:
+        gr.Markdown("""
+        # 🌱 Plant Disease Treatment Assistant
         
-        with gr.Row():
-            with gr.Column():
-                # Image input section
-                gr.Markdown("### Image Analysis")
-                image_input = gr.Image(type="pil", label="Upload Plant Image")
-                image_submit = gr.Button("Analyze Image")
-                
-                # Question input section
-                gr.Markdown("### Ask a Question")
-                question_input = gr.Textbox(
-                    lines=2, 
-                    placeholder="Ask a question like 'How do I treat early blight in tomatoes?'", 
-                    label="Your Question"
-                )
-                question_submit = gr.Button("Get Answer")
+        Upload a plant image or ask a question to get disease identification and treatment information.
+        This tool uses a machine learning model trained on the PlantVillage dataset to identify plant diseases.
+        """)
+        
+        with gr.Tabs():
+            with gr.TabItem("Image Analysis"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        image_input = gr.Image(type="pil", label="Upload Plant Image")
+                        image_submit = gr.Button("Analyze Image", variant="primary")
+                        
+                        gr.Markdown("""
+                        ### Tips for best results:
+                        - Upload a clear, well-lit image
+                        - Focus on affected leaves or plant parts
+                        - Include multiple symptoms if possible
+                        """)
+                        
+                    with gr.Column(scale=2):
+                        analysis_output = gr.Markdown(label="Analysis")
+                        
+                        with gr.Group():
+                            gr.Markdown("### Plant Information")
+                            with gr.Row():
+                                crop_output = gr.Textbox(label="Crop")
+                                disease_output = gr.Textbox(label="Disease")
+                            treatment_output = gr.Textbox(label="Treatment", lines=3)
+                            medicine_output = gr.Textbox(label="Recommended Products", lines=3)
             
-            with gr.Column():
-                # Results display
-                gr.Markdown("### Results")
-                analysis_output = gr.Markdown(label="Analysis")
-                
-                # Structured information display
-                with gr.Group():
-                    gr.Markdown("### Plant Information")
-                    crop_output = gr.Textbox(label="Crop")
-                    disease_output = gr.Textbox(label="Disease")
-                    treatment_output = gr.Textbox(label="Treatment", lines=3)
-                    medicine_output = gr.Textbox(label="Medicine/Chemical Control", lines=3)
+            with gr.TabItem("Q&A"):
+                with gr.Row():
+                    with gr.Column():
+                        question_input = gr.Textbox(
+                            lines=2, 
+                            placeholder="Ask a question like 'How do I treat early blight in tomatoes?'", 
+                            label="Your Question"
+                        )
+                        question_submit = gr.Button("Get Answer", variant="primary")
+                        
+                        gr.Markdown("""
+                        ### Example questions:
+                        - How do I treat early blight in tomatoes?
+                        - What are the symptoms of powdery mildew?
+                        - What chemical controls work on apple scab?
+                        - What causes leaf curl in peach trees?
+                        """)
+                        
+                    with gr.Column():
+                        qa_output = gr.Markdown(label="Answer")
         
         # Set up event handlers
         image_submit.click(
-            fn=lambda img: process_image(img, df),
+            fn=lambda img: process_image(img, model, class_names, transform, df),
             inputs=[image_input],
             outputs=[analysis_output, crop_output, disease_output, treatment_output, medicine_output]
         )
@@ -276,7 +370,7 @@ def main():
         question_submit.click(
             fn=lambda q: (answer_question(q, df), "", "", "", ""),
             inputs=[question_input],
-            outputs=[analysis_output, crop_output, disease_output, treatment_output, medicine_output]
+            outputs=[qa_output, crop_output, disease_output, treatment_output, medicine_output]
         )
         
         # Example questions
@@ -284,14 +378,26 @@ def main():
             examples=[
                 ["How do I treat early blight in tomatoes?"],
                 ["What are the symptoms of powdery mildew?"],
-                ["What chemical controls work on apple scab?"]
+                ["What chemical controls work on apple scab?"],
+                ["How can I prevent late blight in potatoes?"],
+                ["What causes black spots on rose leaves?"]
             ],
             inputs=question_input
         )
+        
+        # Example images
+        gr.Examples(
+            examples=[
+                "example_images/tomato_early_blight.jpg",
+                "example_images/apple_scab.jpg",
+                "example_images/corn_rust.jpg"
+            ],
+            inputs=image_input
+        )
     
     return app
 
 # Launch the app
 if __name__ == "__main__":
     app = main()
-    app.launch()
+    app.launch()