Update app.py

app.py

@@ -1,42 +1,53 @@
+import os
 import gradio as gr
 import tensorflow as tf
 import numpy as np
 import cv2
-import os
 
-# Ensure correct class labels (fix folder naming issue if needed)
-class_labels = ["COVID-19", "Normal", "Pneumonia"]  # Ensure class order is correct
+# Define the model path
+MODEL_PATH = "chest_xray_model.h5"
+
+# Check if the model file exists
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(
+        f"Model file '{MODEL_PATH}' not found. Please upload it to your Hugging Face Space."
+    )
 
 # Load the trained model
-model_path = "chest_xray_model.h5"
-if not os.path.exists(model_path):
-    raise FileNotFoundError(f"Model file '{model_path}' not found. Ensure it's uploaded to the Space.")
+model = tf.keras.models.load_model(MODEL_PATH)
 
-model = tf.keras.models.load_model(model_path)
+# Get class labels from the trained model
+class_labels = ["COVID-19", "NORMAL", "PNEUMONIA"]  # Update if needed
 
-# Preprocessing function for uploaded images
+# Function to preprocess the input image
 def preprocess_image(img):
-    img = cv2.resize(img, (150, 150))  # Resize
+    """Prepares the image for model prediction."""
+    img = cv2.resize(img, (150, 150))  # Resize to match model input shape
     img = img.astype(np.float32) / 255.0  # Normalize pixel values
     img = np.expand_dims(img, axis=0)  # Add batch dimension
     return img
 
-# Prediction function
+# Function to make predictions
 def predict_chest_xray(img):
-    processed_img = preprocess_image(img)
-    prediction = model.predict(processed_img)[0]
-    predicted_class = class_labels[np.argmax(prediction)]
-    confidence = round(100 * np.max(prediction), 2)
-    return f"Prediction: {predicted_class} (Confidence: {confidence}%)"
+    """Runs inference on an uploaded X-ray image."""
+    try:
+        processed_img = preprocess_image(img)
+        prediction = model.predict(processed_img)[0]
+        predicted_class = class_labels[np.argmax(prediction)]
+        confidence = round(100 * np.max(prediction), 2)
+        return f"Prediction: {predicted_class} (Confidence: {confidence}%)"
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Create Gradio UI
+# Create Gradio interface
 interface = gr.Interface(
     fn=predict_chest_xray,
-    inputs=gr.Image(type="numpy"),  # Expect NumPy array
+    inputs=gr.Image(type="numpy"),
     outputs="text",
     title="Chest X-Ray Diagnosis",
-    description="Upload a chest X-ray image to get a diagnosis prediction."
+    description="Upload a chest X-ray image to get a diagnosis prediction.",
 )
 
+# Run the Gradio app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()