Update app.py

app.py

@@ -7,26 +7,26 @@ from PIL import Image
 model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
 feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
 
-# Define class labels and treatment advice
+# Define class labels and treatment advice with a numeric index
 class_labels = {
-    0:"Corn___Common_Rust": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants.",
-    1:"Corn___Gray_Leaf_Spot": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed.",
-    2:"Corn___Healthy": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests.",
-    3:"Corn___Northern_Leaf_Blight": "Remove and destroy infected plant debris, apply fungicides, and rotate crops.",
-    4:"Rice___Brown_Spot": "Use resistant varieties, improve field drainage, and apply fungicides if necessary.",
-    5:"Rice___Healthy": "Maintain proper irrigation, fertilization, and pest control measures.",
-    6:"Rice___Leaf_Blast": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management.",
-    7:"Rice___Neck_Blast": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear.",
-    8:"Wheat___Brown_Rust": "Apply fungicides and practice crop rotation with non-host crops.",
-    9:"Wheat___Healthy": "Continue with good management practices, including proper fertilization and weed control.",
-    10:"Wheat___Yellow_Rust": "Use resistant varieties, apply fungicides, and rotate crops.",
-    11:"Sugarcane__Red_Rot": "Plant resistant varieties and ensure good drainage.",
-    12:"Sugarcane__Healthy": "Maintain healthy soil conditions and proper irrigation.",
-    13:"Sugarcane__Bacterial Blight": "Use disease-free planting material, practice crop rotation, and destroy infected plants."
+    1: {"label": "Corn___Common_Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
+    2: {"label": "Corn___Gray_Leaf_Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
+    3: {"label": "Corn___Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
+    4: {"label": "Corn___Northern_Leaf_Blight", "treatment": "Remove and destroy infected plant debris, apply fungicides, and rotate crops."},
+    5: {"label": "Rice___Brown_Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
+    6: {"label": "Rice___Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
+    7: {"label": "Rice___Leaf_Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
+    8: {"label": "Rice___Neck_Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
+    9: {"label": "Wheat___Brown_Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
+    10: {"label": "Wheat___Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
+    11: {"label": "Wheat___Yellow_Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."},
+    12: {"label": "Sugarcane__Red_Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
+    13: {"label": "Sugarcane__Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
+    14: {"label": "Sugarcane__Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."}
 }
 
 # Mapping label indices to class labels
-labels_list = list(class_labels.keys())
+labels_list = [class_labels[i]["label"] for i in range(1, 15)]
 
 # Inference function
 def predict(image):
@@ -35,10 +35,12 @@ def predict(image):
         outputs = model(**inputs)
         predicted_class_idx = outputs.logits.argmax(-1).item()
     predicted_label = labels_list[predicted_class_idx]
-    treatment_advice = class_labels[predicted_label]
+    
+    # Find corresponding treatment
+    treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
     
     return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
 
 # Create Gradio Interface
 interface = gr.Interface(fn=predict, inputs="image", outputs="text")
-interface.launch()
+interface.launch()