Update app.py

app.py

@@ -1,28 +1,35 @@
 import gradio as gr  # Ensure Gradio is imported
 import torch
 from transformers import ViTForImageClassification, ViTFeatureExtractor
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
-# Define class labels and treatment advice with a numeric index
-# Load model and feature extractor
+
+# Load crop model (ViT)
 model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
 feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
+
+# Load fallback model (ResNet50 for general image classification)
+fallback_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
+fallback_feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
+
+# 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 Corn___Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
+    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 Potato___Early_Blight", "treatment": "Apply fungicides and remove infected plant debris. Practice crop rotation to reduce disease pressure."},
-    6: {"label": "Stage Potato___Healthy", "treatment": "Maintain proper irrigation and fertility practices, and monitor for pests and diseases."},
+    6: {"label": "Stage safe Potato___Healthy", "treatment": "Maintain proper irrigation and fertility practices, and monitor for pests and diseases."},
     7: {"label": "Stage Potato___Late_Blight", "treatment": "Apply fungicides, remove infected plant material, and use resistant potato varieties."},
     8: {"label": "Stage Rice___Brown_Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
-    9: {"label": "Stage Rice___Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
+    9: {"label": "Stage safe Rice___Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
     10: {"label": "Stage Rice___Leaf_Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
     11: {"label": "Stage Rice___Neck_Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
     12: {"label": "Stage Sugarcane__Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."},
-    13: {"label": "Stage Sugarcane__Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
+    13: {"label": "Stage safe Sugarcane__Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
     14: {"label": "Stage Sugarcane__Red_Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
     15: {"label": "Stage Wheat___Brown_Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
-    16: {"label": "Stage Wheat___Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
+    16: {"label": "Stage safe Wheat___Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
     17: {"label": "Stage Wheat___Yellow_Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
 }
 
@@ -31,16 +38,28 @@ labels_list = [class_labels[i]["label"] for i in range(1, 18)]
 
 # Inference function
 def predict(image):
+    # First, use the crop disease model (ViT)
     inputs = feature_extractor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_class_idx = outputs.logits.argmax(-1).item()
-    predicted_label = labels_list[predicted_class_idx]
     
-    # Find corresponding treatment
-    treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
+    # Check if the predicted label corresponds to a crop disease
+    predicted_label = labels_list[predicted_class_idx]
     
-    return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
+    if predicted_class_idx < len(class_labels):  # It's a crop disease
+        treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
+        return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
+    else:
+        # If not a crop disease, use the fallback model (ResNet50) for general object detection
+        inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs_fallback = fallback_model(**inputs_fallback)
+            predicted_class_idx_fallback = outputs_fallback.logits.argmax(-1).item()
+        
+        # Get the fallback prediction label
+        fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
+        return f"Fallback Prediction (Not a Crop): {fallback_label}"
 
 # Create Gradio Interface
 interface = gr.Interface(fn=predict, inputs="image", outputs="text")