Update app.py

app.py

@@ -33,19 +33,21 @@ class_labels = {
 # Mapping label indices to class labels
 labels_list = [class_labels[i]["label"] for i in range(1, 15)]
 
+# Confidence threshold for ViT model
+CONFIDENCE_THRESHOLD = 0.5
+
 # Inference function
-def predict(image, feedback):
+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()
-
-    # Check if the predicted label corresponds to a crop disease
-    predicted_label = labels_list[predicted_class_idx]
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+        confidence = torch.softmax(logits, dim=-1)[0, predicted_class_idx].item()
 
-    # If the feedback is "no", switch to the fallback ResNet-50 model
-    if feedback == "no":
+    # If confidence is below the threshold, directly switch to ResNet50
+    if confidence < CONFIDENCE_THRESHOLD:
         inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
         with torch.no_grad():
             outputs_fallback = fallback_model(**inputs_fallback)
@@ -53,20 +55,20 @@ def predict(image, feedback):
 
         # Get the fallback prediction label
         fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
-        return f"Fallback Prediction (ResNet-50): {fallback_label}"
+        return f"Low confidence in ViT prediction. ResNet-50 Prediction: {fallback_label}"
 
-    # If feedback is "yes", return the initial disease prediction and 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:
-        return f"Disease: {predicted_label}"
+    # If confidence is above the threshold, return the ViT prediction and treatment advice
+    predicted_label = labels_list[predicted_class_idx]
+    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", gr.Radio(["yes", "no"], label="Is it a crop like corn, rice, sugarcane, wheat??")],
-    outputs="text"
+    inputs="image",
+    outputs="text",
+    title="Crop Disease Detection",
+    description="Upload an image of a crop plant to detect diseases. If confidence is low, ResNet-50 will classify the image."
 )
 
 if __name__ == "__main__":