import streamlit as st
import tensorflow as tf
import torch
#import torchvision.transforms as transforms
import numpy as np
import torch.nn as nn
from tensorflow.keras.utils import load_img, img_to_array
from tensorflow.keras.preprocessing import image
from PIL import Image, ImageChops
from torchvision.transforms import transforms
from torchvision import transforms
import torch.nn.functional as F
import pickle
import pandas as pd

classes_p = {'CYST': 0,
    'NORMAL': 1,
    'TUMOR': 2,
    'STONE': 3}
st.sidebar.image("logoMédical.jpg", width=500)
Nom = st.sidebar.button("DONNEES SUR L'AUTEUR")
st.image("logoKeyce.jpg")

if Nom:
    st.title('KEYCE INFORMATIQUE')
    st.title('MASTER 2 IABD')
    st.subheader('TATSA TCHINDA Colince')

m=['CHOISIR UN MODELE DE CLASSIFICATION ICI ',' CLASSIFICATION TENSORFLOW','CLASSIFICATION PYTORCH']
i = st.sidebar.selectbox("Menu", m)
if i== 'CHOISIR UN MODELE DE CLASSIFICATION ICI ':
    st.title('EXAMEN SEMESTRE I, DE DEEP LEARNING CNN')
    st.subheader('Bien vouloir choisir le Framework dans le menu')
    im1 = Image.open("logoTensorFlow.png")
    im2 = Image.open("logoPytorch.jpg")
    taille_image = (400, 300)  # Définir la taille souhaitée
    im1_red = im1.resize(taille_image)
    im2_red = im2.resize(taille_image)
    colonne1, colonne2 = st.columns(2)
    with colonne1:
        st.image(im1_red)

    with colonne2:
        st.image(im2_red)
    
        # Charger les images

elif i == ' CLASSIFICATION TENSORFLOW':
    im1 = Image.open("logoTensorFlow.png")
    taille_image = (800, 200)  # Définir la taille souhaitée
    im1_red = im1.resize(taille_image)
    st.image(im1_red)
    st.title("Classification avec TensorFlow")

    upload_file = st.sidebar.file_uploader('Choisissez une image...',type=['jpg','jpeg','png'])
    generated_pred = st.sidebar.button(' PREDICTION')
    model = tf.keras.models.load_model('TATSA_model_Tensorflow.keras')
    classes_p = {'CYST': 0,
    'NORMAL': 1,
    'TUMOR': 2,
    'STONE': 3}

    if upload_file:
        st.image(upload_file,caption='Image telechargee', use_container_width =True)
        test_image = image.load_img(upload_file,target_size=(64,64))
        image_array = img_to_array(test_image)
        image_array = np.expand_dims(image_array,axis=0)
    else:
        st.markdown('<h3> Attende du scanner ... </h3>', unsafe_allow_html=True)

    if generated_pred:
        predictions = model.predict(image_array)
        classes = np.argmax(predictions[0])
        for key,value in classes_p.items():
            if value == classes:
                st.sidebar.title(f'Je peux dire avec certitude que cette image entre dans la categorie de ➤ {key}')

elif i == 'CLASSIFICATION PYTORCH':
    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('CLASSIFICATION AVEC PYTORCH')

    #Definition du modele

    class Model(nn.Module):
        def __init__(self,dim_output):
            super().__init__()

            self.conv_relu_stack = nn.Sequential(
                nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3),     # (250-3-0)/1 +1  = 248
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2,stride=2))        # (248-2)/2 +1  = 124 
            self.linear_relu_stack = nn.Sequential(
                nn.Linear(in_features=32*124*124, out_features=500),
                nn.ReLU(),
                nn.Linear(in_features=500, out_features=dim_output))
        def forward(self,x):
            x = self.conv_relu_stack(x)
            x = torch.flatten(x,1)
            logits = self.linear_relu_stack(x)
            return logits
        
        
    def load_model():
        model_1 = Model(4)
        model_1.load_state_dict(torch.load('TATSA_classif_py.pth', map_location=torch.device('cpu'), weights_only= True))
        model_1.eval()
        return model_1
        
    model_1 = load_model()

    transform = transforms.Compose([
            transforms.Resize((250, 250)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize for ImageNet
        ])
 
        # Interface utilisateur Streamlit

    uploaded_file = st.sidebar.file_uploader("Choisissez une image...", type=["jpg", "png", "jpeg"])

    generated_pred = st.sidebar.button('PREDICTION')
    if uploaded_file:
        image = Image.open(uploaded_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

    if generated_pred:
            # Prédiction
        with torch.no_grad():
            output = model_1(img_tensor)
            predicted_class = torch.argmax(output, dim=1).item()
            classes_p = {'CYST': 0,'NORMAL': 1,'TUMOR': 2,'STONE': 3}
        for key, value in classes_p.items():
            if value == predicted_class:
                st.title(f'Classification {key}')

    def output_proba(img_tensor):  # Add img_tensor as an argument
        with torch.no_grad():  # Important for inference
            output = model_1(img_tensor) # Get the model's output (logits)
            probabilities = F.softmax(output, dim=1)  # Apply softmax to the output
        return output, probabilities

    if st.sidebar.checkbox("resultat"):
        # Prédiction
        with torch.no_grad():
            output,_ =output_proba(img_tensor)
        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(img_tensor)
            df = pd.DataFrame({"Classe": classes_p.keys(),"Probabilité": probabilities.tolist()[0]})

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