Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 from transformers import ViTForImageClassification, ViTFeatureExtractor
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
-from PIL import Image
 
 # Load models with error handling
 try:
@@ -11,9 +10,9 @@ try:
     fallback_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
     fallback_feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
 except Exception as e:
-    raise gr.Error(f"Failed to load models: {str(e)}. Please check model names and internet connection.")
+    raise gr.Error(f"Model loading failed: {str(e)}")
 
-# Define class labels with treatment advice
+# Class labels and treatments (truncated for brevity)
 class_labels = {
     1: {"label": "Stage Corn Common Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
     2: {"label": "Stage Corn Gray Leaf Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
@@ -31,16 +30,15 @@ class_labels = {
     14: {"label": "Stage Wheat Yellow Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
 }
 
-# Create 0-indexed labels list
 labels_list = [class_labels[i]["label"] for i in range(1, 15)]
-
-# Confidence threshold
 CONFIDENCE_THRESHOLD = 0.5
 
 def predict(image):
     try:
-        # First, try ViT model
-        inputs = feature_extractor(images=image, return_tensors="pt")
+        if not isinstance(image, torch.Tensor):
+            # Convert image to tensor if needed
+            inputs = feature_extractor(images=image, return_tensors="pt")
+        
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
@@ -49,43 +47,57 @@ def predict(image):
             confidence = confidences[0, predicted_class_idx].item()
 
         if confidence < CONFIDENCE_THRESHOLD:
-            # Fallback to ResNet-50
             inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
             with torch.no_grad():
                 outputs_fallback = fallback_model(**inputs_fallback)
-                logits_fallback = outputs_fallback.logits
-                confidences_fallback = torch.softmax(logits_fallback, dim=-1)
-                predicted_class_idx_fallback = logits_fallback.argmax(-1).item()
-                fallback_confidence = confidences_fallback[0, predicted_class_idx_fallback].item()
-
-            fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
-            return (
-                f"Low confidence in ViT model ({confidence * 100:.2f}%).\n"
-                f"ResNet-50 predicts: {fallback_label} ({fallback_confidence * 100:.2f}%).\n\n"
-                "If this doesn't match your input, try another image."
-            )
+                predicted_class_idx_fallback = outputs_fallback.logits.argmax(-1).item()
+                fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
+            return f"Low confidence ({confidence:.2%}). Fallback prediction: {fallback_label}"
 
         predicted_label = labels_list[predicted_class_idx]
-        treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
-        return (
-            f"Disease: {predicted_label} ({confidence * 100:.2f}%)\n\n"
-            f"Treatment Advice: {treatment_advice}"
-        )
+        treatment = class_labels[predicted_class_idx + 1]["treatment"]
+        return f"Disease: {predicted_label}\n\nTreatment: {treatment}"
 
     except Exception as e:
-        return f"Error processing image: {str(e)}. Please try again."
+        return f"Error: {str(e)}"
+
+# Create interface with explicit types
+with gr.Blocks() as demo:
+    gr.Markdown("# 🌱 Crop Disease Detection")
+    gr.Markdown("Upload a crop plant image to detect diseases")
+    
+    with gr.Row():
+        image_input = gr.Image(type="pil")
+        output_text = gr.Textbox()
+    
+    submit_btn = gr.Button("Analyze")
+    submit_btn.click(
+        fn=predict,
+        inputs=image_input,
+        outputs=output_text
+    )
 
-# Gradio Interface
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="🌱 Crop Disease Detection",
-    description="Upload a crop plant image to detect diseases. Uses ViT + ResNet-50 fallback.",
-    examples=[
-        ["corn_rust_example.jpg"],  # Replace with real examples
-        ["wheat_healthy_example.jpg"]
-    ]
-)
+    gr.Examples(
+        examples=[["example_corn.jpg"], ["example_wheat.jpg"]],
+        inputs=image_input
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()
+"""# Define class labels with treatment advice
+class_labels = {
+    1: {"label": "Stage Corn Common Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
+    2: {"label": "Stage Corn Gray Leaf Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
+    3: {"label": "Stage Safe Corn Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
+    4: {"label": "Stage Corn Northern Leaf Blight", "treatment": "Remove and destroy infected plant debris, apply fungicides, and rotate crops."},
+    5: {"label": "Stage Rice Brown Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
+    6: {"label": "Stage Safe Rice Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
+    7: {"label": "Stage Rice Leaf Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
+    8: {"label": "Stage Rice Neck Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
+    9: {"label": "Stage Sugarcane Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."},
+    10: {"label": "Stage Safe Sugarcane Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
+    11: {"label": "Stage Sugarcane Red Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
+    12: {"label": "Stage Wheat Brown Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
+    13: {"label": "Stage Safe Wheat Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
+    14: {"label": "Stage Wheat Yellow Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
+}"""