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

# Define the image size used for training (must match original training)
img_size = 224

# Reconstruct the base model (EfficientNetB0) architecture
base_model = tf.keras.applications.EfficientNetB0(
    include_top=False,
    weights=None, # IMPORTANT: Do not load imagenet weights here directly
    input_shape=(img_size, img_size, 3)
)
# Explicitly build the base model and then load its ImageNet weights from the local file
base_model.build(input_shape=(None, img_size, img_size, 3))
base_model.load_weights('efficientnetb0_notop.h5') # Load pre-downloaded ImageNet weights

# Freeze base_model. The fine-tuned weights will be loaded later.
base_model.trainable = False

# Reconstruct the full model architecture as it was during training
model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(128, activation="relu"),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(4, activation="softmax")
])

# Build the model to create its layers and allocate memory, necessary before loading weights
# Pass an example input shape (None for batch size)
model.build(input_shape=(None, img_size, img_size, 3))

# Load the trained weights into the reconstructed model
model.load_weights('my_model_weights.weights.h5') # This path must match the saved file name

# Class mapping (assuming this is consistent from your training script)
idx_to_class = {0: 'glioma', 1: 'meningioma', 2: 'notumor', 3: 'pituitary'} 
class_labels = list(idx_to_class.values())

def predict_image(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)
    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)

    predictions = model.predict(img_array)[0]
    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.",
    examples=[
        os.path.join("/kaggle/input/brain-tumor-mri-dataset/Testing/notumor/Te-no_0015.jpg"),
        os.path.join("/kaggle/input/brain-tumor-mri-dataset/Testing/glioma/Te-gl_0010.jpg")
    ]
)