Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import torch
 from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
 
@@ -6,22 +7,38 @@ from PIL import Image
 model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
 feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
 
+# Define class labels and treatment advice
+class_labels = {
+    "Corn___Common_Rust": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants.",
+    "Corn___Gray_Leaf_Spot": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed.",
+    "Corn___Healthy": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests.",
+    "Corn___Northern_Leaf_Blight": "Remove and destroy infected plant debris, apply fungicides, and rotate crops.",
+    "Rice___Brown_Spot": "Use resistant varieties, improve field drainage, and apply fungicides if necessary.",
+    "Rice___Healthy": "Maintain proper irrigation, fertilization, and pest control measures.",
+    "Rice___Leaf_Blast": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management.",
+    "Rice___Neck_Blast": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear.",
+    "Wheat___Brown_Rust": "Apply fungicides and practice crop rotation with non-host crops.",
+    "Wheat___Healthy": "Continue with good management practices, including proper fertilization and weed control.",
+    "Wheat___Yellow_Rust": "Use resistant varieties, apply fungicides, and rotate crops.",
+    "Sugarcane__Red_Rot": "Plant resistant varieties and ensure good drainage.",
+    "Sugarcane__Healthy": "Maintain healthy soil conditions and proper irrigation.",
+    "Sugarcane__Bacterial Blight": "Use disease-free planting material, practice crop rotation, and destroy infected plants."
+}
+
+# Mapping label indices to class labels
+labels_list = list(class_labels.keys())
+
+# Inference function
 def predict(image):
-    # Preprocess image
     inputs = feature_extractor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    class_labels = model.config.id2label  # Assuming you have a class label mapping
-    return class_labels[predicted_class_idx]
-
-# Gradio interface
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Crop Disease Classifier",
-    description="Upload an image of a crop leaf to classify the disease."
-)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_class_idx = outputs.logits.argmax(-1).item()
+    predicted_label = labels_list[predicted_class_idx]
+    treatment_advice = class_labels[predicted_label]
+    
+    return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
 
+# Create Gradio Interface
+interface = gr.Interface(fn=predict, inputs="image", outputs="text")
 interface.launch()