Update app.py

app.py

@@ -45,6 +45,11 @@ if image_display is not None:
     if st.button("Predict"):
         infer = model.signatures["serving_default"]
         output = infer(tf.constant(image_array))
-        prediction = output[list(output.keys())[0]].numpy()
-        predicted_class = labels[np.argmax(prediction)]
-        st.success(f"Prediction: {predicted_class}")
+        prediction = output[list(output.keys())[0]].numpy()[0]
+        predicted_index = np.argmax(prediction)
+        predicted_class = labels[predicted_index]
+        confidence = prediction[predicted_index]
+        st.success(f"Prediction: {predicted_class.upper()} (Confidence: {confidence:.2f})")
+        st.subheader("Class probabilities")
+        for label, prob in zip(labels, prediction):
+            st.write(f"{label}: {prob:.2f}")