Update app.py

app.py

@@ -11,8 +11,11 @@ classifier = VisionClassifierInference(
 
 # Define a Gradio interface
 def classify_image(image_file):
+    """Classify an image using a pre-trained ViT model."""
     label = classifier.predict(img_path=image_file.name)
-    return f"Predicted class: {label}"
+    # Add a confidence score to the output
+    confidence = classifier.predict_proba(img_path=image_file.name)[0][label]
+    return f"Predicted class: {label} (confidence: {confidence:.2f})"
 
 iface = gr.Interface(
     fn=classify_image,