import streamlit as st
import torch
from torchvision import transforms
from PIL import Image
import torch.nn.functional as F
import pickle
import pandas as pd
me = ['Acceuil','Analyse Medicale']
p = st.sidebar.selectbox('menu', me)

@st.cache_resource  # Pour éviter de recharger le modècdle à chaque interaction
def charger_modele_pytorch():
    modele = torch.load('TATSA_Model1_TransfLr_py.pth', map_location=device)  # Charge le modèle
    #modele.eval()  # Important : mettez le modèle en mode évaluation
    return modele
#chemin_modele = charger_modele_pytorch()
#chemin_modele = st.text_input("Chemin vers le modèle (.pth)", "pytorch1.pth")

if torch.cuda.is_available():
    device = st.selectbox("Utiliser le GPU ou le CPU ?", ["GPU", "CPU"], index=0)
    device = torch.device("cuda" if device == "GPU" else "cpu")
else:
    #st.write("Aucun GPU détecté. Utilisation du CPU.")
    device = torch.device("cpu")

modele_charge = charger_modele_pytorch()

    # Transformation des images (Doit être la même que celle utilisée à l'entraînement)
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Modifier selon la taille utilisée à l'entraînement
    transforms.ToTensor(),
    transforms.Normalize((0.5),0.5)])
if p == 'Acceuil':
    st.title('EXAMEN SEMESTRE I, DE TRANSFERT LEARNING')
    st.subheader('MASTER 2 IABD')
    st.subheader('TATSA TCHINDA Colince')
    
elif p=='Analyse Medicale':
    st.title('TRANSFERT LEARNING')
    upload_file = st.sidebar.file_uploader('Choisissez une image',type=['jpg','jpeg','png'])
    
    if upload_file:
        image = Image.open(upload_file).convert("RGB")
        st.image(image, caption="Image chargée", use_container_width=True)

        # Prétraitement de l'image
        img_tensor = transform(image).unsqueeze(0)  # Ajout d'une dimension batch
        #bouton1 = st.sidebar.button('Resultat')
        #bouton2 = st.sidebar.buttonRadio('Probabilites')
        classes_p = {'HEALTHY': 0,'BRAIN_TUMOR': 1}

        def output_proba():
            output = modele_charge(img_tensor)
            probabilities = F.softmax(output, dim=1)
            return output, probabilities

        if st.sidebar.checkbox("resultat"):
            # Prédiction
            with torch.no_grad():
                output,_ =output_proba()
            predicted_class = torch.argmax(output, dim=1).item()
            
            for key, value in classes_p.items():
                if value == predicted_class:
                    st.title(f'Categorie de ➤ {key}')
            if st.sidebar.checkbox("prob"):
                        
                _,probabilities = output_proba()
                df = pd.DataFrame({"Classe": classes_p.keys(),"Probabilité": probabilities.tolist()[0]})

                st.write("Probabilités :")
                st.dataframe(df)