from PIL import Image, ExifTags
import streamlit as st
from transformers import pipeline

import torch
# Assurez-vous que models.py est correctement défini et dans votre PATH
from models import vgg19  
from torchvision import transforms


# Initialisation des pipelines Hugging Face
caption_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
emotion_pipeline = pipeline("image-classification", model="RickyIG/emotion_face_image_classification_v3")

# Chemin vers votre modèle local et initialisation
model_path = "model_sh_B.pth"
device = torch.device('cpu')  # Changez pour 'cuda' si vous utilisez un GPU

# Charger le modèle vgg19
model = vgg19()
model.to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()


def predict_count(image):
    # Prétraitement de l'image
    trans = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    img_tensor = trans(image)
    inp = img_tensor.unsqueeze(0).to(device)

    with torch.no_grad():
        outputs, _ = model(inp)

    count = torch.sum(outputs).item()
    return int(count)

def generate_caption_emotion_and_people_count(image):
    # Générer une légende 1 ere etape
    caption_result = caption_pipeline(image)
    caption = caption_result[0]["generated_text"]

    # Classification des émotions 2 eme etape
    emotion_result = emotion_pipeline(image)
    emotions = ", ".join([f"{res['label']}: {res['score']:.2f}" for res in emotion_result])

    # Comptage des personnes dans la photo 3 eme etaoe
    count = predict_count(image)

    # Combinaison des résultats
    # combined_result = f"Caption: {caption}\nEmotions: {emotions}\nNumber of People: {count}"
    # return combined_result

    return caption, emotions, count
def main():
    # Interface Streamlit
    st.title("Analyse d'Affluence Événementielle via la Détection d'Objets")

    # Présentation des objectifs du projet
    st.write("Ce projet vise à :")
    st.markdown("""
        - Compter le nombre de participants.
        - Analyser leurs émotions.
        - Générer une description contextuelle de l'événement.
    """)

    # Upload d'image
    uploaded_image = st.file_uploader("Choisissez une image...", type=["jpg", "jpeg", "png"])
    if uploaded_image is not None:




        image = Image.open(uploaded_image)

        # Obtenir et afficher la taille de l'image
        width, height = image.size
        if width>1200 or height > 1200:
            try:
                exif = image._getexif()
                if exif is not None:
                    orientation_key = next((key for key, value in ExifTags.TAGS.items() if value == 'Orientation'), None)
                    if orientation_key is not None and orientation_key in exif:
                        orientation = exif[orientation_key]
                        if orientation == 3:
                            image = image.rotate(180, expand=True)
                        elif orientation == 6:
                            image = image.rotate(270, expand=True)
                        elif orientation == 8:
                            image = image.rotate(90, expand=True)
                        # Affichage de l'image uploadée
                        width, height = image.size
                        st.write(f"Largeur: {width} pixels, Hauteur: {height} pixels")
                        st.image(image, caption='Image Uploadée', use_column_width=True)
            except Exception as e:
                st.write(f"Erreur lors de la correction de l'orientation: {e}")
                pass
            image = image.resize((224, 224))
        else:
            # Affichage de l'image uploadée
            width, height = image.size
            st.write(f"Largeur: {width} pixels, Hauteur: {height} pixels")
            st.image(image, caption='Image Uploadée', use_column_width=True)
        
        
        


        # Placeholder pour le bouton ou le message de chargement
        button_placeholder = st.empty()

        # Si le bouton est cliqué
        if button_placeholder.button('Analyser l\'image'):
            # Affichage du spinner et du message pendant le traitement
            with st.spinner('En cours d\'exécution...'):
                caption, emotions, count = generate_caption_emotion_and_people_count(image)
                
                # Remplacement du spinner par les résultats une fois le traitement terminé
                with st.expander("Voir les résultats !!"):
                    st.write(f"**Légende**: {caption}")
                    st.write(f"**Émotions**: {emotions}")
                    st.write(f"**Nombre de personnes**: {count}")
            
            # Optionnellement, effacez le placeholder ou affichez un message différent après l'exécution
            button_placeholder.empty()

if __name__ == '__main__':
    main()