Update app.py

app.py

@@ -6,7 +6,7 @@ import warnings
 import sys
 import os
 import contextlib
-from transformers import ViTForImageClassification
+from transformers import ViTForImageClassification, pipeline
 
 # Suppress warnings related to the model weights initialization
 warnings.filterwarnings("ignore", category=UserWarning, message=".*weights.*")
@@ -39,6 +39,9 @@ transform = transforms.Compose([
 # Load the class names (disease types)
 class_names = ['BacterialBlights', 'Healthy', 'Mosaic', 'RedRot', 'Rust', 'Yellow']
 
+# 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")
+
 # Function to predict disease type from an image
 def predict_disease(image):
     # Apply transformations to the image
@@ -51,7 +54,11 @@ def predict_disease(image):
     
     # Get the predicted label
     predicted_label = class_names[predicted_class.item()]
-    
+
+    # Generate AI response based on the detected disease
+    prompt = f"You are an expert in sugarcane farming. The detected disease is '{predicted_label}'. Provide advice for the farmer."
+    ai_response = ai_pipeline(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']
+
     # Create a styled HTML output
     output_message = f"""
     <div style='font-size: 18px; color: #4CAF50; font-weight: bold;'>
@@ -62,7 +69,7 @@ def predict_disease(image):
     if predicted_label != "Healthy":
         output_message += f"""
         <p style='font-size: 16px; color: #757575;'>
-            This indicates the presence of <strong>{predicted_label}</strong>. Please take immediate action to prevent further spread.
+            <strong>AI Response:</strong> {ai_response}
         </p>
         """
     else:
@@ -90,4 +97,6 @@ demo_app = gr.Interface(
     theme="huggingface"
 )
 
+demo_app.launch(debug=True)
+
 demo_app.launch(debug=True)