import streamlit as st
import numpy as np 
from tensorflow.keras.models import load_model 
from tensorflow.keras.preprocessing import image 
import matplotlib.pyplot as plt 

# Disable the PyplotGlobalUseWarning 
st.set_option('deprecation.showPyplotGlobalUse', False)

# Load the model
model = load_model('brain_tumor_model.keras')

# Class mappings
class_mappings = {
    'Glioma': 0,
    'Meningioma': 1,
    'Notumor': 2,
    'Pituitary': 3
}
inv_class_mappings = {v: k for k, v in class_mappings.items()}
class_names = list(class_mappings.keys())

# Load the true and predicted labels
true_labels = np.load('true_labels.npy') # Model trained on the Brain MRI Dataset
predicted_labels = np.load('predicted_labels.npy') # Output response after Training/Testing TF CNN Model

# Note0: The One Hot Encode predictions are leveraging the TensorFlow Convolutional Neural Network Model
# Note1: The One Hot Encode predictions are deriving from the actual Brain MRI Dataset images

# Function to load and preprocess an image
def load_and_preprocess_image(image_path, image_shape=(168, 168)):
    img = image.load_img(image_path, target_size=image_shape, color_mode='grayscale') 
    img_array = image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array 

# Function to display a row of images with predictions
def display_images_and_predictions(image_paths, predictions, true_labels, figsize=(20,5)):
    fig, axes = plt.subplots(1, len(image_paths), figsize=figsize)
    for i, (image_path, prediction, true_label) in enumerate(zip(image_paths, predictions, true_labels)):
        ax = axes[i]
        img_array = load_and_preprocess_image(image_path)
        img_array = np.squeeze(img_array)
        ax.imshow(img_array, cmap='gray')
        title_color = 'green' if prediction == true_label else 'red'
        ax.set_title(f'True Label: {true_label}\nPred: {prediction}', color=title_color)
        ax.axis('off')
    st.pyplot(fig)


# Image paths
image_paths = [
    'Te-gl_0057.jpg',
    'Te-me_0057.jpg',
    'Te-no_0057.jpg',
    'Te-pi_0057.jpg'
    
]

# True labels for images
true_labels = [
    'Glioma', 'Meningioma', 'Notumor', 'Pituitary', # Original tumor types
]

# Load and preprocess images, then make predictions
images = [load_and_preprocess_image(path) for path in image_paths] 
predictions = [model.predict(image) for image in images] 

# Determine the predicted labels
predicted_labels = [inv_class_mappings[np.argmax(one_hot)] for one_hot in predictions] 

# Create Streamlit app title
st.markdown("<hi style='text-align: center; color: navy;'>Brain Tumor One Hot Encode TF Model Version 3</h1>", unsafe_allow_html=True)

# Output the predictions
st.write(f'Class Mappings: {class_mappings}')
st.write("\nGlioma Image Prediction:", np.round(predictions[0], 3)[0])
st.write("Meningioma Image Prediction:", np.round(predictions[1], 3)[0])
st.write("Notumor Image Prediction:", np.round(predictions[2], 3)[0])
st.write("Pituitary Image Prediction:", np.round(predictions[3], 3)[0])

# Display images with predictions
display_images_and_predictions(image_paths, predicted_labels, true_labels)