Update app.py

app.py

@@ -2,188 +2,90 @@ import streamlit as st
 import numpy as np
 import cv2
 from PIL import Image
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
 from io import BytesIO
-#from ultralytics import YOLO
 import os
-import tempfile
-import base64
-import requests
-from datetime import datetime
-import google.generativeai as genai  # Import Gemini API
-from tensorflow.keras.models import load_model
-from vit import create_vit_classifier
-# Load the model
-vit_classifier = load_model('models/vit_updated.weights.h5')  # Replace with your model file path
-# Configuring Google Gemini API
-GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY")
-genai.configure(api_key=GEMINI_API_KEY)
+import google.generativeai as genai
+from vit_model import VisionTransformer
 
-# Loading YOLO model for crop disease detection
-#yolo_model = YOLO("models/best.pt")
+# 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
 
-# Initializing conversation history if not set
-if "conversation_history" not in st.session_state:
-    st.session_state.conversation_history = {}
-
-# Function to preprocess images
-def preprocess_image(image, target_size=(256, 256)):
-    """Resize image for AI models."""
-    image = Image.fromarray(image)
-    image = image.resize(target_size)
-    return image
+# Configure Google Gemini API
+GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY")
+genai.configure(api_key=GEMINI_API_KEY)
 
-# Generate response from Gemini AI with history
-def generate_gemini_response(disease_list, user_context="", conversation_history=None):
-    """Generate a structured diagnosis using Gemini API, considering conversation history."""
+# Load the model
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = VisionTransformer(img_size=128, patch_size=8, num_classes=len(CLASS_NAMES), embed_dim=768, depth=8, num_heads=12, mlp_dim=2048, dropout=0.1)
+model.load_state_dict(torch.load("custom_vit.pth", map_location=DEVICE))
+model.to(DEVICE)
+model.eval()
+
+# Function to preprocess images for ViT
+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
+    ])
+    return transform(image).unsqueeze(0).to(DEVICE)
+
+# Function for inference using ViT
+def vit_inference(image):
+    """Predicts crop disease using the custom ViT model."""
+    input_tensor = preprocess_image(image)
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        predicted_class = torch.argmax(probabilities).item()
+    return CLASS_NAMES[predicted_class], probabilities[predicted_class].item()
+
+# Generate response from Gemini AI
+def generate_gemini_response(disease_name):
+    """Generate a structured diagnosis using Gemini API."""
     try:
         model = genai.GenerativeModel("gemini-1.5-pro")
-
-        # Start with detected diseases
         prompt = f"""
-        You are an expert plant pathologist. The detected crop diseases is: {predicted_labels}.
+        You are an expert plant pathologist. The detected crop disease is: {disease_name}.
         
-        User's context or question: {user_context if user_context else "Provide a general analysis"}
+        Provide a structured analysis including:
+        - Pathogen details
+        - Severity level
+        - Symptoms
+        - Economic impact
+        - Treatment options (short-term and long-term)
+        - Prevention strategies
         """
-
-        # Add past conversation history for better continuity
-        if conversation_history:
-            history_text = "\n\nPrevious conversation:\n"
-            for entry in conversation_history:
-                history_text += f"- User: {entry['question']}\n- AI: {entry['response']}\n"
-            prompt += history_text
-
-        # Ask Gemini for a structured diagnosis
-        prompt += """
-        For each detected disease, provide a structured analysis following this format:
-        1. Disease Name: [Name]
-       - Pathogen: [Causative organism]
-       - Severity Level: [Based on visual symptoms]
-       - Key Symptoms:
-         * [Symptom 1]
-         * [Symptom 2]
-       - Economic Impact:
-         * [Brief description of potential crop losses]
-       - Treatment Options:
-         * Immediate actions: [Short-term solutions]
-         * Long-term management: [Preventive measures]
-       - Environmental Conditions:
-         * Favorable conditions for disease development
-         * Risk factors
-        2. Recommendations:
-       - Immediate Steps:
-         * [Action items for immediate control]
-       - Prevention Strategy:
-         * [Long-term prevention measures]
-       - Monitoring Protocol:
-         * [What to watch for]"""
-
-
         response = model.generate_content(prompt)
         return response.text if response else "No response from Gemini."
     except Exception as e:
         return f"Error connecting to Gemini API: {str(e)}"
 
-# Performing inference using YOLO
-def inference(image):
-    """Detect crop diseases in the given image with confidence filtering."""
-    predictions = vit_classifier.predict(image)
-    predicted_labels = np.argmax(predictions, axis=1)
-
-    return predicted_labels
-
-
+# Initialize Streamlit app
+st.title("🌱 AI-Powered Crop Disease Detection")
 
-
-# Initialize Streamlit UI
-st.title("AI-Powered Crop Disease Detection & Diagnosis System")
-
-# Sidebar settings
-with st.sidebar:
-    st.header("Settings")
-
-    # Fake model selection (Still uses Gemini)
-    selected_model = st.selectbox("Choose Model", ["Gemini", "GPT-4", "Claude", "Llama 3", "Mistral"], help="This app always uses Gemini.")
-
-    confidence_threshold = st.slider("Detection Confidence Threshold", 0.0, 1.0, 0.4)
-    
-  
-
-    if st.button("Clear Conversation History"):
-        st.session_state.conversation_history = {}
-        st.success("Conversation history cleared!")
-
-
-# User context input with example prompts
-st.subheader("Provide Initial Context or Ask a Question")
-
-# Generalized example prompts for easier input
-example_prompts = {
-    "Select an example...": "",
-    "General Plant Health Issue": "My plant leaves are wilting and turning yellow. Is this a disease or a nutrient deficiency?",
-    "Leaf Spots and Discoloration": "I see dark spots on my crop leaves. Could this be a fungal or bacterial infection?",
-    "Leaves Drying or Curling": "The leaves on my plants are curling and drying up. What could be causing this?",
-    "Pest or Disease?": "I noticed tiny insects on my plants along with some leaf damage. Could this be a pest problem or a disease?",
-    "Overwatering or Root Rot?": "My plant leaves are turning brown and mushy. Is this due to overwatering or a root infection?",
-    "Poor Crop Growth": "My crops are growing very slowly and seem weak. Could this be due to soil problems or disease?",
-    "Weather and Disease Connection": "It has been raining a lot, and now my plants have mold. Could the weather be causing a fungal disease?",
-    "Regional Disease Concern": "I'm in a humid area and my crops often get infected. What are common diseases for this climate?",
-}
-
-# Dropdown menu for selecting an example
-selected_example = st.selectbox("Choose an example to auto-fill:", list(example_prompts.keys()))
-
-# Auto-fill the text area when an example is selected
-user_context = st.text_area(
-    "Enter details, symptoms, or a question about your plant condition.",
-    value=example_prompts[selected_example] if selected_example != "Select an example..." else "",
-    placeholder="Example: My plant leaves are turning yellow and wilting. Is this a disease or a nutrient issue?"
-)
-
-# Upload an image
+# Upload image
 uploaded_file = st.file_uploader("📤 Upload a plant image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    file_id = uploaded_file.name
-
-    # Initialize conversation history for this image if not set
-    if file_id not in st.session_state.conversation_history:
-        st.session_state.conversation_history[file_id] = []
-
-    # Convert file to image
-    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-    img = cv2.imdecode(file_bytes, 1)
-
-    # Perform inference
-    predicted_labels = inference(img)
-
-    # Display processed image with detected diseases
-    st.image(img, caption="🔍 Detected Diseases", use_column_width=True)
-
-  
-    st.write(f"✅ **High Confidence Diseases Detected:** {predicted_labels}")
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
     
-
-
-        # AI-generated diagnosis from Gemini
+    # Run ViT inference
+    with st.spinner("Analyzing with Vision Transformer... 🔍"):
+        predicted_class, confidence = vit_inference(image)
+    
+    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)
     st.subheader("📋 AI Diagnosis")
-    with st.spinner("Generating diagnosis... 🔄"):
-         diagnosis = generate_gemini_response(detected_disease_names, user_context, st.session_state.conversation_history[file_id])
-
-        # Save response to history
-    st.session_state.conversation_history[file_id].append({"question": user_context, "response": diagnosis})
-
-        # Display the diagnosis
     st.write(diagnosis)
 
-        # Show past conversation history
-    if st.session_state.conversation_history[file_id]:
-        st.subheader("🗂️ Conversation History")
-        for i, entry in enumerate(st.session_state.conversation_history[file_id]):
-            with st.expander(f"Q{i+1}: {entry['question'][:50]}..."):
-                st.write("**User:**", entry["question"])
-                st.write("**AI:**", entry["response"]) 
-       
-
 # Instructions for users
 st.markdown("""
 ---