import gradio as gr
import tensorflow as tf
from PIL import Image, ImageOps
import numpy as np

# Load the model
model = tf.keras.models.load_model("pneumonia_cnn_model.h5")

# Prediction function
def predict(image):
    try:
        # 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
        pred_prob = model.predict(img_array)[0][0]  # Extract float value

        # Interpret prediction
        if pred_prob >= 0.5:
            return f"Pneumonia detected (confidence: {pred_prob:.2f})"
        else:
            return f"No pneumonia detected (confidence: {pred_prob:.2f})"

    except Exception as e:
        return f"Error during prediction: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,  # Function to call on image input
    inputs=gr.Image(type="pil", label="Upload Chest X-ray Image"),
    outputs="text",  # Output is the prediction result (text)
    live=True  # Optional: set to False if you don't want to update results live
)

# Launch the app
iface.launch()