Update app.py

app.py

@@ -8,10 +8,11 @@ import os
 import contextlib
 from transformers import ViTForImageClassification, pipeline
 
-# Suppress warnings
+# Suppress warnings related to the model weights initialization
 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:
@@ -22,12 +23,16 @@ def suppress_stdout():
         finally:
             sys.stdout = old_stdout
 
-# Load the saved model and suppress the warnings
+# Load the sugarcane disease model
 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)),
@@ -35,12 +40,9 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Load the class names
+# Load the class names (disease types)
 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.",
@@ -51,10 +53,31 @@ knowledge_base = {
     'Healthy': "The sugarcane crop is healthy. Continue regular monitoring and good agronomic practices."
 }
 
-# Update the predict_disease function to handle non-sugarcane images
+# 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
 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)  # Add batch dimension
+    img_tensor = transform(image).unsqueeze(0)
 
     # Make prediction
     with torch.no_grad():
@@ -62,65 +85,37 @@ 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.8  # 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
-    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
+    detailed_response = knowledge_base.get(predicted_label, "No additional information available.")
+
+    # Create 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>
     """
-    
-    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>
-        """
-    
+    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",
-    examples=EXAMPLES,
-    live=True,
-    theme="huggingface"
+    live=True
 )
 
 demo_app.launch(debug=True)
 
+