Update app.py

app.py

@@ -8,11 +8,10 @@ import os
 import contextlib
 from transformers import ViTForImageClassification, pipeline
 
-# Suppress warnings related to the model weights initialization
+# Suppress warnings
 warnings.filterwarnings("ignore", category=UserWarning, message=".*weights.*")
 warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
 
-# Suppress output for copying files and verbose model initialization messages
 @contextlib.contextmanager
 def suppress_stdout():
     with open(os.devnull, 'w') as devnull:
@@ -23,16 +22,12 @@ def suppress_stdout():
         finally:
             sys.stdout = old_stdout
 
-# Load the sugarcane disease model
+# Load the saved model and suppress the warnings
 with suppress_stdout():
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=6)
     model.load_state_dict(torch.load('vit_sugarcane_disease_detection.pth', map_location=torch.device('cpu')))
     model.eval()
 
-# Load a general-purpose classifier (e.g., MobileNetV2)
-general_model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)
-general_model.eval()
-
 # Define the same transformation used during training
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -40,9 +35,12 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Load the class names (disease types)
+# Load the class names
 class_names = ['BacterialBlights', 'Healthy', 'Mosaic', 'RedRot', 'Rust', 'Yellow']
 
+# Load AI response generator
+ai_pipeline = pipeline("text-generation", model="gpt2", tokenizer="gpt2")
+
 # Knowledge base for sugarcane diseases
 knowledge_base = {
     'BacterialBlights': "Bacterial blights cause water-soaked lesions on leaves, leading to yellowing and withering. To manage, apply copper-based fungicides and ensure proper drainage.",
@@ -53,31 +51,10 @@ knowledge_base = {
     'Healthy': "The sugarcane crop is healthy. Continue regular monitoring and good agronomic practices."
 }
 
-# Function to check if the image is plant-related
-def is_plant_image(image):
-    general_transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-    ])
-    img_tensor = general_transform(image).unsqueeze(0)
-    with torch.no_grad():
-        outputs = general_model(img_tensor)
-        _, predicted_class = torch.max(outputs, 1)
-    # Check if the predicted class corresponds to plant-like images
-    plant_related_classes = range(20, 25)  # Replace with specific classes for plants
-    return predicted_class.item() in plant_related_classes
-
-# Predict disease or detect non-sugarcane images
+# Update the predict_disease function to handle non-sugarcane images
 def predict_disease(image):
-    if not is_plant_image(image):
-        return """
-        <div style='font-size: 18px; color: #FF5722; font-weight: bold;'>
-            The uploaded image is not related to sugarcane. Please upload a sugarcane image.
-        </div>
-        """
-    
     # Apply transformations to the image
-    img_tensor = transform(image).unsqueeze(0)
+    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
 
     # Make prediction
     with torch.no_grad():
@@ -85,35 +62,64 @@ def predict_disease(image):
         probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
         max_prob, predicted_class = torch.max(probabilities, 1)
     
+    # Confidence threshold for non-sugarcane detection
+    confidence_threshold = 0.6  # Adjust based on experimentation
+
+    # Check if the confidence is below the threshold
+    if max_prob.item() < confidence_threshold:
+        return """
+        <div style='font-size: 18px; color: #FF5722; font-weight: bold;'>
+            The uploaded image does not belong to the sugarcane dataset.
+        </div>
+        """
+
     # Get the predicted label
     predicted_label = class_names[predicted_class.item()]
 
     # Retrieve response from knowledge base
-    detailed_response = knowledge_base.get(predicted_label, "No additional information available.")
-
-    # Create styled HTML output
+    if predicted_label in knowledge_base:
+        detailed_response = knowledge_base[predicted_label]
+    else:
+        # Fallback to AI-generated response
+        prompt = f"The detected sugarcane disease is '{predicted_label}'. Provide detailed advice for managing this condition."
+        detailed_response = ai_pipeline(prompt, max_length=100, num_return_sequences=1, truncation=True)[0]['generated_text']
+
+    # Create a styled HTML output
     output_message = f"""
     <div style='font-size: 18px; color: #4CAF50; font-weight: bold;'>
         Detected Disease: <span style='color: #FF5722;'>{predicted_label}</span>
     </div>
     """
-    output_message += f"""
-    <p style='font-size: 16px; color: #757575;'>
-        {detailed_response}
-    </p>
-    """
+    
+    if predicted_label != "Healthy":
+        output_message += f"""
+        <p style='font-size: 16px; color: #757575;'>
+            {detailed_response}
+        </p>
+        """
+    else:
+        output_message += f"""
+        <p style='font-size: 16px; color: #757575;'>
+            {detailed_response}
+        </p>
+        """
+    
     return output_message
 
 # Create Gradio interface
 inputs = gr.Image(type="pil")
 outputs = gr.HTML()
 
+EXAMPLES = ["img1.jpeg", "redrot2.jpg", "rust1.jpg", "healthy2.jpeg"]
+
 demo_app = gr.Interface(
     fn=predict_disease,
     inputs=inputs,
     outputs=outputs,
     title="Sugarcane Disease Detection",
-    live=True
+    examples=EXAMPLES,
+    live=True,
+    theme="huggingface"
 )
 
 demo_app.launch(debug=True)