Update app.py

app.py

@@ -3,27 +3,31 @@ import numpy as np
 import cv2
 from PIL import Image
 from io import BytesIO
-from ultralytics import YOLO
+#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 vit_classifier
 
+# Load the model
+model = load_model('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)
 
 # Loading YOLO model for crop disease detection
-yolo_model = YOLO("models/best.pt")
+#yolo_model = YOLO("models/best.pt")
 
 # 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=(224, 224)):
+def preprocess_image(image, target_size=(256, 256)):
     """Resize image for AI models."""
     image = Image.fromarray(image)
     image = image.resize(target_size)
@@ -81,24 +85,12 @@ def generate_gemini_response(disease_list, user_context="", conversation_history
         return f"Error connecting to Gemini API: {str(e)}"
 
 # Performing inference using YOLO
-def inference(image, conf_threshold=0.5):
+def inference(image):
     """Detect crop diseases in the given image with confidence filtering."""
-    results = yolo_model(image, conf=0.4)  # Adjusted confidence threshold for detection
-    infer = np.zeros(image.shape, dtype=np.uint8)
-    detected_classes = []
-    class_names = {}
-    confidence_scores = []
+    predictions = vit_classifier.predict(image)
+    predicted_labels = np.argmax(predictions, axis=1)
 
-    for r in results:
-        infer = r.plot()
-        class_names = r.names
-        for i, cls in enumerate(r.boxes.cls.tolist()):
-            confidence = r.boxes.conf[i].item()  # Get confidence score
-            if confidence >= conf_threshold:  # Only consider high-confidence detections
-                detected_classes.append(cls)
-                confidence_scores.append(confidence)
-
-    return infer, detected_classes, class_names, confidence_scores
+    return predicted_labels
 
 
 
@@ -163,21 +155,13 @@ if uploaded_file:
     img = cv2.imdecode(file_bytes, 1)
 
     # Perform inference
-    processed_image, detected_classes, class_names, confidence_scores = inference(img)
+    predicted_labels = inference(img)
 
     # Display processed image with detected diseases
-    st.image(processed_image, caption="🔍 Detected Diseases", use_column_width=True)
-
-    if detected_classes:
-    # Convert detected class indexes to names
-        detected_disease_names = [
-            f"{class_names[cls]} ({confidence_scores[i]:.2f})"
-            for i, cls in enumerate(detected_classes)
-    ]
-
-    # Show only the most confident detections
-    if detected_disease_names:
-        st.write(f"✅ **High Confidence Diseases Detected:** {', '.join(detected_disease_names)}")
+    st.image(img, caption="🔍 Detected Diseases", use_column_width=True)
+
+  
+    st.write(f"✅ **High Confidence Diseases Detected:** {predicted_labels)}")