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 = ['Accueil','Analyse Medicale']
st.sidebar.image("logoMédical.jpg", width=500)
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, weights_only=False)  # 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(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])])
if p == 'Accueil':
    st.image("logoKeyce.jpg")
    st.title('KEYCE INFORMATIQUE')
    st.title('EXAMEN SEMESTRE I, DE TRANSFERT LEARNING')
    st.subheader('MASTER 2 IABD')
    st.subheader('TATSA TCHINDA Colince')
    
elif p=='Analyse Medicale':
    st.image("logoKeyce.jpg")
    im2 = Image.open("logoPytorch.jpg")
    taille_image = (800, 200)  # Définir la taille souhaitée
    im2_red = im2.resize(taille_image)
    st.image(im2_red)
    st.title('TRANSFERT LEARNING AVEC PYTORCH')
    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("Probabilités"):
                        
                _,probabilities = output_proba()
                df = pd.DataFrame({"Classe": classes_p.keys(),"Probabilité": probabilities.tolist()[0]})

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