Update app.py

app.py

@@ -33,11 +33,14 @@ def predict(image):
     
     # Perform inference
     predictions = model.predict(input_tensor)
+
+    predicted_probabilities = np.squeeze(predictions)
+    predicted_class = np.argmax(predicted_probabilities)
         
     # Extract the predicted class label with the highest probability
-    prediction = np.argmax(predictions[0])
+    # prediction = np.argmax(predictions[0])
 
-    if(prediction > 0.5):
+    if(predicted_class > 0.5):
         return "Pneumonia"
     else:
         return "Normal"