Update app.py

app.py

@@ -32,15 +32,17 @@ def predict(image):
     # Get the predicted label
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    
+    confidence_score = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_class_idx].item()
+
     # Get the label name
     predicted_label = id2label[str(predicted_class_idx)]
-    
-    return predicted_label
+
+    # Return the predicted label and confidence score
+    return predicted_label, f"Confidence: {confidence_score:.2f}"
 
 # Create the Gradio interface
 image = gr.Image(type="pil")
 label = gr.Label(num_top_classes=3)
 
 gr.Interface(fn=predict, inputs=image, outputs=label, title="Citrus Disease Classification",
-             description="Upload an image of a citrus leaf or fruit to classify its disease.").launch()
+             description="Upload an image of a citrus leaf to classify its disease.").launch()