# importing the libraries and dependencies needed for creating the UI and supporting the deep learning models used in the project
import streamlit as st
import tensorflow as tf
import random
from PIL import Image
from tensorflow import keras
import numpy as np
import os

import warnings

warnings.filterwarnings("ignore")
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

st.set_page_config(
            page_title="ChestAI - Pneumonia Detection",
            page_icon="🫁",
            initial_sidebar_state="expanded"
        )

hide_streamlit_style = """
	<style>
  #MainMenu {visibility: hidden;}
	footer {visibility: hidden;}
  </style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)


def prediction_cls(prediction):
    for key, clss in class_names.items():  # create a dictionary of the output classes
        if np.argmax(prediction) == clss:  # check the class
            return key


with st.sidebar:
    st.title("👋 Welcome to ChestAI")
            st.markdown("""
                ### About
                ChestAI uses advanced deep learning to detect pneumonia in chest X-rays.
                
                ### How to use
                1. Upload a chest X-ray image (JPG/PNG)
                2. Wait for the analysis
                3. View the results and confidence score
                
                ### Note
                This tool is for educational purposes only. Always consult healthcare professionals for medical advice.
            """)
st.set_option("deprecation.showfileUploaderEncoding", False)


@st.cache_resource()
def load_model():
    from huggingface_hub import from_pretrained_keras

    keras.utils.set_random_seed(42)
    model = from_pretrained_keras("ryefoxlime/PneumoniaDetection")
    return model


with st.spinner("Model is being loaded.."):
    model = load_model()

file = st.file_uploader(" ", type=["jpg", "png"])


def import_and_predict(image_data, model):
    img_array = keras.preprocessing.image.img_to_array(image_data)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = img_array/255

    predictions = model.predict(img_array)
    return predictions


if file is None:
    st.text("Please upload an image file")
else:
    image = keras.preprocessing.image.load_img(file, target_size=(224, 224), color_mode='rgb')
    st.image(image, caption="Uploaded Image.", use_column_width=True)
    predictions = import_and_predict(image, model)
    print(predictions)
    class_names = [
        "Normal",
        "PNEUMONIA",
    ]

    string = "Detected Disease : " + class_names[np.argmax(predictions)]
    if class_names[np.argmax(predictions)] == "Normal":
        st.balloons()
        st.success(string)

    elif class_names[np.argmax(predictions)] == "PNEUMONIA":
        st.warning(string)