Update app.py

app.py

@@ -1,76 +1,179 @@
 import streamlit as st
-import torch
+import requests
 from PIL import Image
-from transformers import ViTForImageClassification, ViTImageProcessor
+import torch
+from torchvision import transforms
+import os
+from groq import Groq
+
+# Initialize Groq client
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+# Set up the app title and layout
+st.set_page_config(page_title="Leaves Disease Detection", layout="wide")
+st.title("🌿 Leaves Disease Detection")
+st.write("Upload an image of a plant leaf to check for diseases and get treatment recommendations.")
 
-# Load model and processor
+# Load the plant disease classification model
 @st.cache_resource
 def load_model():
-    processor = ViTImageProcessor.from_pretrained('nateraw/vit-base-patch16-224-plant-village')
-    model = ViTForImageClassification.from_pretrained('nateraw/vit-base-patch16-224-plant-village')
-    return processor, model
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path='plant_disease_model.pt', force_reload=True)
+    model.eval()
+    return model
 
-processor, model = load_model()
+model = load_model()
 
-# Class labels mapping (simplified example)
-class_labels = {
-    0: "Apple Scab",
-    1: "Apple Black Rot",
-    2: "Apple Cedar Rust",
-    3: "Apple Healthy",
-    4: "Cherry Healthy",
-    5: "Cherry Powdery Mildew",
-    # Add all 38 classes from PlantVillage dataset
-    # Complete mapping available at: https://github.com/nateraw/plant-village-classifier
-}
+# Image preprocessing
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    return transform(image).unsqueeze(0)
 
-# Treatment information
-treatment_info = {
-    "Apple Scab": {
-        "summary": "Fungal disease affecting apples",
-        "treatment": [
-            "Remove and destroy infected leaves",
-            "Apply fungicides containing fenarimol or myclobutanil",
-            "Prune trees to improve air circulation"
-        ]
-    },
-    "Apple Black Rot": {
-        "summary": "Fungal disease causing fruit rot",
-        "treatment": [
-            "Remove mummified fruits",
-            "Apply copper-based fungicides",
-            "Practice proper sanitation in orchard"
+# Disease classification
+def classify_disease(image):
+    try:
+        # Preprocess the image
+        img_tensor = preprocess_image(image)
+        
+        # Perform inference
+        with torch.no_grad():
+            outputs = model(img_tensor)
+        
+        # Get the predicted class
+        _, predicted = torch.max(outputs, 1)
+        class_idx = predicted.item()
+        
+        # Map class index to disease name (this should match your model's classes)
+        disease_classes = [
+            "Healthy", "Apple Scab", "Apple Black Rot", "Apple Cedar Rust", 
+            "Cherry Powdery Mildew", "Corn Gray Leaf Spot", "Corn Common Rust", 
+            "Grape Black Rot", "Grape Esca", "Grape Leaf Blight", 
+            "Orange Huanglongbing", "Peach Bacterial Spot", "Pepper Bacterial Spot", 
+            "Potato Early Blight", "Potato Late Blight", "Raspberry Healthy", 
+            "Soybean Healthy", "Squash Powdery Mildew", "Strawberry Leaf Scorch", 
+            "Tomato Bacterial Spot", "Tomato Early Blight", "Tomato Late Blight", 
+            "Tomato Leaf Mold", "Tomato Septoria Leaf Spot", "Tomato Spider Mites", 
+            "Tomato Target Spot", "Tomato Yellow Leaf Curl Virus", "Tomato Mosaic Virus"
         ]
-    },
-    # Add treatments for other diseases
-}
-
-st.title("🌱 Plant Disease Diagnosis & Treatment")
+        
+        disease_name = disease_classes[class_idx]
+        return disease_name
+    except Exception as e:
+        st.error(f"Error during classification: {str(e)}")
+        return "Unknown"
 
-uploaded_file = st.file_uploader("Upload a leaf image...", type=["jpg", "jpeg", "png"])
+# Get disease description and treatment from Groq API
+def get_disease_info(disease_name):
+    try:
+        if disease_name.lower() == "healthy":
+            return {
+                "description": "The plant appears to be healthy with no visible signs of disease.",
+                "treatment": "No treatment needed. Continue with regular plant care practices.",
+                "prevention": "Maintain good growing conditions, proper watering, and regular monitoring."
+            }
+        
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a plant pathologist assistant. Provide accurate information about plant diseases, their symptoms, treatments, and prevention methods."
+                },
+                {
+                    "role": "user",
+                    "content": f"Provide a detailed description of the plant disease {disease_name}, its symptoms, recommended treatments (including organic and chemical options), and prevention methods. Format the response with clear sections for Description, Treatment, and Prevention."
+                }
+            ],
+            model="mixtral-8x7b-32768",
+            temperature=0.3,
+            max_tokens=1024
+        )
+        
+        response = chat_completion.choices[0].message.content
+        return parse_groq_response(response)
+    except Exception as e:
+        st.error(f"Error fetching disease information: {str(e)}")
+        return {
+            "description": "Information not available.",
+            "treatment": "Please consult a local agricultural expert.",
+            "prevention": "Regular monitoring and good cultural practices are recommended."
+        }
 
-if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Leaf Image", use_column_width=True)
+def parse_groq_response(response):
+    # Simple parsing of the Groq response
+    sections = {
+        "description": "",
+        "treatment": "",
+        "prevention": ""
+    }
     
-    # Preprocess and predict
-    with st.spinner("Analyzing..."):
-        inputs = processor(images=image, return_tensors="pt")
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predicted_class_idx = logits.argmax(-1).item()
-        
-    disease = class_labels.get(predicted_class_idx, "Unknown Disease")
+    current_section = None
+    for line in response.split('\n'):
+        line_lower = line.lower()
+        if "description" in line_lower:
+            current_section = "description"
+            sections[current_section] += line + "\n"
+        elif "treatment" in line_lower:
+            current_section = "treatment"
+            sections[current_section] += line + "\n"
+        elif "prevention" in line_lower:
+            current_section = "prevention"
+            sections[current_section] += line + "\n"
+        elif current_section:
+            sections[current_section] += line + "\n"
     
-    st.subheader("Diagnosis Results")
-    st.success(f"Detected Disease: **{disease}**")
+    return sections
+
+# Main application
+def main():
+    uploaded_file = st.file_uploader("Upload a leaf image", type=["jpg", "jpeg", "png"])
     
-    if disease in treatment_info:
-        st.subheader("Disease Summary")
-        st.write(treatment_info[disease]["summary"])
+    if uploaded_file is not None:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Leaf Image", use_column_width=True)
         
-        st.subheader("Recommended Treatment")
-        for step, treatment in enumerate(treatment_info[disease]["treatment"], 1):
-            st.write(f"{step}. {treatment}")
-    else:
-        st.warning("Treatment information not available for this disease.")
+        if st.button("Predict Disease"):
+            with st.spinner("Analyzing the leaf..."):
+                # Classify the disease
+                disease_name = classify_disease(image)
+                
+                # Get disease information
+                disease_info = get_disease_info(disease_name)
+                
+                # Display results
+                st.subheader("🌱 Analysis Results")
+                col1, col2 = st.columns(2)
+                
+                with col1:
+                    st.markdown(f"**Status:** {'Healthy' if disease_name.lower() == 'healthy' else 'Diseased'}")
+                    st.markdown(f"**Detected Condition:** {disease_name}")
+                
+                with col2:
+                    if disease_name.lower() != "healthy":
+                        st.warning("⚠️ Disease Detected")
+                    else:
+                        st.success("✅ Healthy Plant")
+                
+                st.subheader("📝 Detailed Information")
+                
+                with st.expander("Description"):
+                    st.write(disease_info["description"])
+                
+                if disease_name.lower() != "healthy":
+                    with st.expander("Recommended Treatment"):
+                        st.write(disease_info["treatment"])
+                    
+                    with st.expander("Prevention Methods"):
+                        st.write(disease_info["prevention"])
+                
+                st.subheader("🌍 Real-world Summary")
+                if disease_name.lower() == "healthy":
+                    st.write(f"The analysis indicates that the plant leaf appears healthy with no signs of disease. Healthy plants typically have vibrant color, uniform texture, and no visible spots or discoloration. Maintaining proper growing conditions, including adequate sunlight, water, and nutrients, will help keep the plant healthy.")
+                else:
+                    st.write(f"The analysis detected {disease_name}, a common plant disease that can affect plant health and productivity. {disease_info['description'].split('.')[0]}. Early detection and proper treatment are crucial to prevent the spread of the disease and protect other plants in the vicinity. Following the recommended treatment and prevention methods can help restore plant health and prevent future outbreaks.")
+
+if __name__ == "__main__":
+    main()