samixFr's picture
update image.size and rotate
c65aa6a verified
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()