import streamlit as st
import pickle
from Prediction import  Prediction
from huggingface_hub import hf_hub_download


# resnet_model = pickle.load(open('models/ResNet01.pkl','rb'))
# cnn_model = pickle.load(open('models/CNNModel2.pkl','rb'))
# inc_model = pickle.load(open('models/Inception01.pkl','rb'))
# resnet_model = TFAutoModel.from_pretrained("yashpat85/ResNet01")

REPO_ID1 = "yashpat85/ResNet01"
MODEL_DIR1 = hf_hub_download(repo_id=REPO_ID1, filename="ResNet01.pkl", repo_type="model")

REPO_ID2= "yashpat85/Inception01"
MODEL_DIR2 = hf_hub_download(repo_id=REPO_ID2, filename="Inception01.pkl", repo_type="model")

REPO_ID3 = "yashpat85/CNNModel2"
MODEL_DIR3 = hf_hub_download(repo_id=REPO_ID3, filename="CNNModel2.pkl", repo_type="model")


resnet_model = pickle.load(open(MODEL_DIR1,'rb'))
cnn_model = pickle.load(open(MODEL_DIR3,'rb'))
inc_model = pickle.load(open(MODEL_DIR2,'rb'))

def show_error_popup(message):
    st.error(message, icon="🚨")

st.set_page_config(layout="wide")

st.title('Kidney Disease Classification using CNN')
st.markdown('By 22DCS079 & 22DCS085')

st.header('Add Ct Scan Image')

uploaded_file = st.file_uploader("Choose a ct scan image", type=["jpg", "png", "jpeg"])

st.header("Available Models")
option = st.selectbox(
    "Available Models",
    ("ResNet", "CNN","InceptionNet"),
)

pm = Prediction()

col1, col2= st.columns(2)

if uploaded_file is not None:
    with col1:
        image_data = uploaded_file.read()
        st.image(image_data, caption="Uploaded Image")
    with col2:
        if option=="CNN":
            p = pm.predict_image(cnn_model, image_data)
        elif option=="ResNet":
            p = pm.predict_image(resnet_model, image_data)
        elif option=="InceptionNet":
            p = pm.predict_image(inc_model, image_data)
        else:
            p = "Other Models are still under training due to overfitting"

        print(p)

        if p=='Normal':
            st.markdown("""
            <style>
            .big-font {
                display: flex;
                align-items:center;
                justify-content: center;
                font-size:50px !important;
                color:green;
                height: 50vh;
            }
            </style>
            """, unsafe_allow_html=True)

            st.markdown(f'<div class="big-font">{p}</div>', unsafe_allow_html=True)
        else:
            st.markdown("""
            <style>
            .big-font {
                display: flex;
                align-items:center;
                justify-content: center;
                font-size:50px !important;
                color:red;
                height: 50vh;
            }
            </style>
            """, unsafe_allow_html=True)

            st.markdown(f'<div class="big-font">{p}</div>', unsafe_allow_html=True)
else:
    show_error_popup("Please Upload Image...")