Update app.py

app.py

@@ -4,39 +4,36 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 import os
+from dotenv import load_dotenv
 from groq import Groq
 
-# Set up the app
+# Load environment variables
+load_dotenv()
+
+# Streamlit UI setup
 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.")
 
 # Initialize Groq client
 try:
-    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+    api_key = os.getenv("GROQ_API_KEY")
+    client = Groq(api_key=api_key)
 except Exception as e:
     st.error(f"Failed to initialize Groq client: {str(e)}")
     client = None
 
-# Simple CNN model for plant disease classification
+# Simple CNN model (dummy architecture)
 class PlantDiseaseModel(nn.Module):
-    def __init__(self, num_classes=38):
+    def __init__(self, num_classes=28):
         super(PlantDiseaseModel, self).__init__()
         self.features = nn.Sequential(
-            nn.Conv2d(3, 32, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(32, 64, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-            nn.Conv2d(64, 128, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),
+            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),
+            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),
         )
         self.classifier = nn.Sequential(
-            nn.Linear(128 * 32 * 32, 512),
-            nn.ReLU(),
-            nn.Dropout(0.5),
+            nn.Linear(128 * 32 * 32, 512), nn.ReLU(), nn.Dropout(0.5),
             nn.Linear(512, num_classes)
         )
     
@@ -46,7 +43,7 @@ class PlantDiseaseModel(nn.Module):
         x = self.classifier(x)
         return x
 
-# Load model (dummy implementation)
+# Cache the model
 @st.cache_resource
 def load_model():
     model = PlantDiseaseModel()
@@ -55,16 +52,30 @@ def load_model():
 
 model = load_model()
 
-# Image preprocessing
+# Preprocess image
 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]),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                             [0.229, 0.224, 0.225]),
     ])
     return transform(image).unsqueeze(0)
 
-# Disease classification
+# Dummy disease 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"
+]
+
+# Classify the image
 def classify_disease(image):
     try:
         img_tensor = preprocess_image(image)
@@ -72,99 +83,69 @@ def classify_disease(image):
             outputs = model(img_tensor)
         _, predicted = torch.max(outputs, 1)
         class_idx = predicted.item()
-        
-        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"
-        ]
-        
-        disease_name = disease_classes[class_idx % len(disease_classes)]
-        return disease_name
+        return disease_classes[class_idx % len(disease_classes)]
     except Exception as e:
         st.error(f"Error during classification: {str(e)}")
         return "Unknown"
 
-# Get disease info from Groq API
+# Fetch info from Groq API
 def get_disease_info(disease_name):
     if not client:
         return {
-            "description": "API connection not available. Please check your Groq API key.",
-            "treatment": "",
-            "prevention": ""
+            "description": "API connection not available. Please check your GROQ_API_KEY.",
         }
-    
     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."
+                "description": "The plant appears to be healthy. No treatment needed.",
             }
-        
+
         response = client.chat.completions.create(
             messages=[
                 {"role": "system", "content": "You are a plant pathologist assistant."},
-                {"role": "user", "content": f"Describe {disease_name} in plants with symptoms, treatment, and prevention."}
+                {"role": "user", "content": f"Describe {disease_name} in plants including symptoms, treatment, and prevention."}
             ],
             model="mixtral-8x7b-32768",
             temperature=0.3,
             max_tokens=1024
         )
-        
         return {"description": response.choices[0].message.content}
     except Exception as e:
         st.error(f"Error fetching disease information: {str(e)}")
         return {
-            "description": "Information not available. Please consult an expert.",
-            "treatment": "",
-            "prevention": ""
+            "description": "Unable to fetch disease info. Please try again later.",
         }
 
-# Main app
+# Main app function
 def main():
     uploaded_file = st.file_uploader("Upload a leaf image", type=["jpg", "jpeg", "png"])
     
-    if uploaded_file is not None:
+    if uploaded_file:
         try:
-            image = Image.open(uploaded_file)
+            image = Image.open(uploaded_file).convert("RGB")
             st.image(image, caption="Uploaded Leaf Image", use_column_width=True)
-            
-            if st.button("Predict Disease"):
+
+            if st.button("🔍 Predict Disease"):
                 with st.spinner("Analyzing the leaf..."):
                     disease_name = classify_disease(image)
-                    disease_info = get_disease_info(disease_name)
-                    
-                    st.subheader("🌱 Analysis Results")
+                    info = get_disease_info(disease_name)
+
+                    st.subheader("🔬 Prediction Results")
                     col1, col2 = st.columns(2)
-                    
                     with col1:
                         status = "Healthy" if disease_name.lower() == "healthy" else "Diseased"
                         st.markdown(f"**Status:** {status}")
-                        st.markdown(f"**Detected Condition:** {disease_name}")
-                    
+                        st.markdown(f"**Detected Disease:** {disease_name}")
                     with col2:
-                        if disease_name.lower() != "healthy":
-                            st.warning("⚠️ Disease Detected")
+                        if disease_name.lower() == "healthy":
+                            st.success("✅ Plant is Healthy")
                         else:
-                            st.success("✅ Healthy Plant")
-                    
-                    st.subheader("📝 Detailed Information")
-                    st.write(disease_info["description"])
-                    
-                    st.subheader("🌍 Real-world Summary")
-                    if disease_name.lower() == "healthy":
-                        st.write("The analysis indicates a healthy plant leaf with no signs of disease.")
-                    else:
-                        st.write(f"The analysis detected {disease_name}. Early detection and proper treatment are crucial.")
+                            st.warning("⚠️ Disease Detected")
+
+                    st.subheader("📋 Detailed Information")
+                    st.write(info["description"])
         except Exception as e:
             st.error(f"Error processing image: {str(e)}")
 
 if __name__ == "__main__":
-    main()
+    main()