import numpy as np
import pandas as pd
import altair as alt
from PIL import Image
import streamlit as st
import tensorflow as tf
from itertools import chain

col1, col2, col3 = st.columns(3)

with col1:
    st.write(' ')

with col2:
    st.image("unilogo.png", width=300)
    st.markdown("<h6 style='text-align: center; color: grey;'>University of Technology and Applied Sciences, Muscat</h6>",
            unsafe_allow_html=True)
    st.markdown("<h1 style='text-align: center; color: grey;'>X-RAY Classsifier</h1>",
                unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>Project Supervisor : DEEVYANKAR AGARWAL</h6>",
            unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>MALIKA SULTAN ALNAABI</h6>",
            unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>MAATHIR YAHYA ABDULLAH AL HINAI</h6>",
            unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>MARYAM MOHAMMED HILAL AL-HASHAMI</h6>",
            unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: grey;'>SHIFA NASSER MOHAMMED AL AAMRI</h6>",
            unsafe_allow_html=True)

with col3:
    st.write(' ')


file_upload = st.file_uploader("Upload the X-RAY Scan ", type=[
                               "png", "jpg", "jpeg", "heic", "tiff", "bmp", "raw"], accept_multiple_files=False)


class_dict = {"0": 'covid19', "1": 'normal', "2": 'pneumonia'}

model = tf.keras.models.load_model('fmodel.h5')


def preprocess_image(image_path):

    image = Image.open(image_path)

    image = image.resize((224, 224))
    image_array = tf.keras.preprocessing.image.img_to_array(image)
    image_array /= 255.0
    image_array = tf.expand_dims(image_array, 0)

    return image_array


if file_upload:
    
    with open(file_upload.name, "wb") as f:
        f.write(file_upload.getbuffer())

    saved_path = f"{file_upload.name}"

    preprocessed_image = preprocess_image(saved_path)
    pred_prob = model.predict(preprocessed_image)
    predicted_class_index = np.argmax(pred_prob, axis=1).item()
    predicted_class = class_dict[f"{predicted_class_index}"]

    st.write("Predicted Class: ", predicted_class)

    print(predicted_class_index)
    print(predicted_class)
    print(pred_prob)

    graph_input = list(chain.from_iterable(pred_prob.tolist()))
    "Plot depicting Class Probabilities"
    source = pd.DataFrame({
        'Class probability': graph_input,
        'Class': ["COVID-19", "NORMAL", "PNEUMONIA"]
    })

    bar_chart = alt.Chart(source).mark_bar().encode(
        y='Class probability:Q',
        x='Class:O',
    )

    element5 = st.altair_chart(bar_chart, use_container_width=True)