Update app.py

app.py

@@ -19,14 +19,14 @@ class_labels = {
     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."},
+    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."},
+    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."},
+    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."}
 }
 
@@ -46,7 +46,7 @@ def predict(image):
         predicted_class_idx = logits.argmax(-1).item()
         confidence = torch.softmax(logits, dim=-1)[0, predicted_class_idx].item()
 
-    # If confidence is below the threshold, directly switch to ResNet50
+    # If confidence is below the threshold, use the fallback model
     if confidence < CONFIDENCE_THRESHOLD:
         inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
         with torch.no_grad():
@@ -55,7 +55,8 @@ def predict(image):
 
         # Get the fallback prediction label
         fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
-        return f"Low confidence in ViT prediction. ResNet-50 Prediction: {fallback_label}"
+
+        return f"Low confidence in ViT model. ResNet-50 suggests: {fallback_label}\n\nIf this does not match your input, please try another image."
 
     # If confidence is above the threshold, return the ViT prediction and treatment advice
     predicted_label = labels_list[predicted_class_idx]