Update app.py

app.py

@@ -9,23 +9,26 @@ model = tf.keras.models.load_model("pneumonia_cnn_model.h5")
 # Prediction function
 def predict(image):
     try:
-        # Convert to grayscale
-        img = ImageOps.grayscale(image)
-        
-        # Resize image to match the model's expected input size (299x299)
-        img = img.resize((299, 299))  # Resize to 299x299
-        
-        # Convert to numpy array and normalize the pixel values
-        img_array = np.array(img).reshape(1, 299, 299, 1) / 255.0  # Reshape to (1, 299, 299, 1)
-        
+        # Resize image to match the model's expected input size
+        img = image.resize((299, 299))
+
+        # Handle grayscale vs RGB based on model input shape
+        if model.input_shape[-1] == 1:  # Model trained on grayscale
+            img = ImageOps.grayscale(img)
+            img_array = np.array(img).reshape(1, 299, 299, 1) / 255.0
+        else:  # Model trained on RGB
+            img_array = np.array(img) / 255.0
+            img_array = img_array.reshape(1, 299, 299, 3)
+
         # Make prediction
-        prediction = model.predict(img_array)
-        
+        pred_prob = model.predict(img_array)[0][0]  # Extract float value
+
         # Interpret prediction
-        if prediction >= 0.5:
-            return "Pneumonia detected"
+        if pred_prob >= 0.5:
+            return f"Pneumonia detected (confidence: {pred_prob:.2f})"
         else:
-            return "No pneumonia detected"
+            return f"No pneumonia detected (confidence: {pred_prob:.2f})"
+
     except Exception as e:
         return f"Error during prediction: {str(e)}"