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

# Load the trained model
model = tf.keras.models.load_model('brain_tumor_classifier.keras')

# Get class names from the training data generator (assuming 'train' is still in scope from previous execution)
# If 'train' is not in scope, you would need to define class_indices manually or reload data generators.
# For this example, let's assume 'train.class_indices' is available or define a placeholder.

# If `train` is not available, uncomment and modify the line below based on your actual classes:
idx_to_class = {0: 'glioma', 1: 'meningioma', 2: 'notumor', 3: 'pituitary'}

# Using the `idx_to_class` from previous execution
# If `idx_to_class` is not defined, please refer to the notebook output from the prediction cell.
class_labels = list(idx_to_class.values())

# Define the image size used for training
img_size = 224

def predict_image(image):
    # Preprocess the image
    img = Image.fromarray(image)
    img = img.resize((img_size, img_size))
    img_array = np.array(img)
    img_array = np.expand_dims(img_array, axis=0) # Add batch dimension

    # Apply the same preprocessing function as during training
    # (EfficientNet's preprocess_input function was used)
    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)

    # Make prediction
    predictions = model.predict(img_array)[0]

    # Get predicted class and confidence
    predicted_class_idx = np.argmax(predictions)
    predicted_class_label = class_labels[predicted_class_idx]
    confidence = predictions[predicted_class_idx] * 100

    return predicted_class_label, f"{confidence:.2f}%"

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy", label="Upload MRI Scan"),
    outputs=[
        gr.Textbox(label="Predicted Class"),
        gr.Textbox(label="Confidence")
    ],
    title="Brain Tumor MRI Classification",
    description="Upload an MRI scan to get a prediction for brain tumor type and confidence.",
)

# Launch the interface
iface.launch(debug=True)