Update app.py

app.py

@@ -36,32 +36,43 @@ labels_list = [class_labels[i]["label"] for i in range(1, 15)]
 # Confidence threshold for ViT model
 CONFIDENCE_THRESHOLD = 0.5
 
-# Inference function
+# Inference function with fuzzy confidence
 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)
         logits = outputs.logits
+        confidences = torch.softmax(logits, dim=-1)
         predicted_class_idx = logits.argmax(-1).item()
-        confidence = torch.softmax(logits, dim=-1)[0, predicted_class_idx].item()
+        confidence = confidences[0, predicted_class_idx].item()
 
     # 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():
             outputs_fallback = fallback_model(**inputs_fallback)
-            predicted_class_idx_fallback = outputs_fallback.logits.argmax(-1).item()
+            logits_fallback = outputs_fallback.logits
+            confidences_fallback = torch.softmax(logits_fallback, dim=-1)
+            predicted_class_idx_fallback = logits_fallback.argmax(-1).item()
+            fallback_confidence = confidences_fallback[0, predicted_class_idx_fallback].item()
 
         # Get the fallback prediction label
         fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
 
-        return f"Low confidence in ViT model. ResNet-50 suggests: {fallback_label}\n\nIf this does not match your input, please try another image."
+        return (
+            f"Low confidence in ViT model ({confidence * 100:.2f}%).\n"
+            f"ResNet-50 predicts: {fallback_label} ({fallback_confidence * 100:.2f}%).\n\n"
+            "If 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]
     treatment_advice = class_labels[predicted_class_idx + 1]["treatment"]
-    return f"Disease: {predicted_label}\n\nTreatment Advice: {treatment_advice}"
+    return (
+        f"Disease: {predicted_label} ({confidence * 100:.2f}%)\n\n"
+        f"Treatment Advice: {treatment_advice}"
+    )
 
 # Create Gradio Interface
 interface = gr.Interface(