Update app.py

app.py

@@ -23,10 +23,21 @@ def predict_pneumonia(image):
 
     if prediction >= threshold:
         pneumonia_prediction = 1
+        class_name = "Pneumonia"
     else:
         pneumonia_prediction = 0
+        class_name = "Normal"
+
+    # Return the probability of each class
+    class_probabilities = model.predict(img_array)[0]
+
+    return {
+        "Pneumonia Prediction": pneumonia_prediction,
+        "Class Name": class_name,
+        "Pneumonia Probability": class_probabilities[0],
+        "Normal Probability": class_probabilities[1]
+    }
 
-    return pneumonia_prediction
 
 inputs = gr.inputs.Image(shape=(180, 180))
 outputs = gr.outputs.Label(num_top_classes=2)