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:
@@ -35,13 +36,13 @@ 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
+# Load AI response generator (using a local GPT pipeline or OpenAI's GPT-3/4 API)
 ai_pipeline = pipeline("text-generation", model="gpt2", tokenizer="gpt2")
 
-# Knowledge base for sugarcane diseases
+# Knowledge base for sugarcane diseases (example data from the website)
 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.",
     'Mosaic': "Mosaic disease results in streaked and mottled leaves, reducing photosynthesis. Use disease-resistant varieties and control aphids to prevent spread.",
@@ -51,7 +52,7 @@ 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
+# Update the predict_disease function
 def predict_disease(image):
     # Apply transformations to the image
     img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
@@ -59,20 +60,8 @@ def predict_disease(image):
     # Make prediction
     with torch.no_grad():
         outputs = model(img_tensor)
-        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
-        max_prob, predicted_class = torch.max(probabilities, 1)
+        _, predicted_class = torch.max(outputs.logits, 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()]
 
@@ -108,7 +97,7 @@ def predict_disease(image):
 
 # Create Gradio interface
 inputs = gr.Image(type="pil")
-outputs = gr.HTML()
+outputs = gr.HTML()  # Use HTML output for styled text
 
 EXAMPLES = ["img1.jpeg", "redrot2.jpg", "rust1.jpg", "healthy2.jpeg"]
 
@@ -124,4 +113,3 @@ demo_app = gr.Interface(
 
 demo_app.launch(debug=True)
 
-