Update app.py

app.py

@@ -10,8 +10,8 @@ import os
 import google.generativeai as genai
 from vit_model import VisionTransformer
 
-# Load class names (should match training labels)
-CLASS_NAMES =  ['Pepper__bell___Bacterial_spot', 'Pepper__bell___healthy', 'Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy', 'Tomato_Bacterial_spot', 'Tomato_Early_blight', 'Tomato_Late_blight', 'Tomato_Leaf_Mold', 'Tomato_Septoria_leaf_spot', 'Tomato_Spider_mites_Two_spotted_spider_mite', 'Tomato__Target_Spot', 'Tomato__Tomato_YellowLeaf__Curl_Virus', 'Tomato__Tomato_mosaic_virus', 'Tomato_healthy']  # Update with actual class names
+# Load class names
+CLASS_NAMES =  ['Pepper__bell___Bacterial_spot', 'Pepper__bell___healthy', 'Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy', 'Tomato_Bacterial_spot', 'Tomato_Early_blight', 'Tomato_Late_blight', 'Tomato_Leaf_Mold', 'Tomato_Septoria_leaf_spot', 'Tomato_Spider_mites_Two_spotted_spider_mite', 'Tomato__Target_Spot', 'Tomato__Tomato_YellowLeaf__Curl_Virus', 'Tomato__Tomato_mosaic_virus', 'Tomato_healthy']
 
 # Configure Google Gemini API
 GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY")
@@ -29,7 +29,7 @@ def preprocess_image(image, target_size=(128, 128)):
     transform = transforms.Compose([
         transforms.Resize(target_size),
         transforms.ToTensor(),
-        transforms.Normalize([0.5], [0.5])  # Normalize to match training
+        transforms.Normalize([0.5], [0.5])
     ])
     return transform(image).unsqueeze(0).to(DEVICE)
 
@@ -44,13 +44,15 @@ def vit_inference(image):
     return CLASS_NAMES[predicted_class], probabilities[predicted_class].item()
 
 # Generate response from Gemini AI
-def generate_gemini_response(disease_name):
+def generate_gemini_response(disease_name, user_context):
     """Generate a structured diagnosis using Gemini API."""
     try:
         model = genai.GenerativeModel("gemini-1.5-pro")
         prompt = f"""
         You are an expert plant pathologist. The detected crop disease is: {disease_name}.
         
+        User-provided context: {user_context}
+        
         Provide a structured analysis including:
         - Pathogen details
         - Severity level
@@ -65,10 +67,11 @@ def generate_gemini_response(disease_name):
         return f"Error connecting to Gemini API: {str(e)}"
 
 # Initialize Streamlit app
-st.title("🌱 AI-Powered Crop Disease Detection")
+st.title("AI-Powered Crop Disease Detection with User Context")
 
 # Upload image
 uploaded_file = st.file_uploader("📤 Upload a plant image", type=["jpg", "jpeg", "png"])
+user_context = st.text_area("📝 Additional Context (Optional)", "Describe symptoms, location, or any observations...")
 
 if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
@@ -78,11 +81,11 @@ if uploaded_file:
     with st.spinner("Analyzing with Vision Transformer... 🔍"):
         predicted_class, confidence = vit_inference(image)
     
-    st.write(f"✅ **Detected Disease:** {predicted_class} (Confidence: {confidence:.2f})")
+    st.write(f"**Detected Disease:** {predicted_class} (Confidence: {confidence:.2f})")
     
     # Connect to Gemini AI for diagnosis
     with st.spinner("Generating diagnosis with Gemini AI... 💡"):
-        diagnosis = generate_gemini_response(predicted_class)
+        diagnosis = generate_gemini_response(predicted_class, user_context)
     st.subheader("📋 AI Diagnosis")
     st.write(diagnosis)
 
@@ -94,5 +97,4 @@ st.markdown("""
 2. Provide context (optional) about symptoms or concerns.
 3. The system detects the disease using AI.
 4. Gemini generates a diagnosis with symptoms and treatments.
-5. Ask follow-up questions, and the AI will remember previous responses.
 """)