Update app.py

app.py

@@ -37,7 +37,7 @@ def generate_text_with_gemini(api_key, prompt):
     return response.status_code, response_json
 
 # Fonction de traitement de l'image et d'encadrement des personnes détectées
-def process_image(image):
+def process_image(image, confidence_threshold=0.7):
     processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
     model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
     
@@ -47,7 +47,7 @@ def process_image(image):
     
     # Post-traitement des résultats pour obtenir les boîtes englobantes et les labels
     target_sizes = torch.tensor([image.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=confidence_threshold)[0]
 
     # Convertir l'image pour dessiner
     draw_image = image.copy()
@@ -56,44 +56,55 @@ def process_image(image):
     # Compter le nombre de personnes détectées
     persons = 0
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        if model.config.id2label[label.item()] == 'person':
+        if score >= confidence_threshold and model.config.id2label[label.item()] == 'person':
             persons += 1
             box = [round(i, 2) for i in box.tolist()]
             # Dessiner le rectangle rouge autour de la personne détectée
             draw.rectangle(box, outline="red", width=3)
     
-    expected_number_of_persons = 10  # Nombre attendu de personnes
-    absents = expected_number_of_persons - persons
-
-    # Assurez-vous que le nombre d'absents ne soit pas négatif
-    if absents < 0:
-        absents = 0
-
-    # Générer le contenu de l'email avec Gemini
-    prompt = f"Notifier de l'absence: Le nombre de personnes présentes est de {persons}. Le nombre d'absents est de {absents}."
-    gemini_api_key = "AIzaSyCJVUr-76I4TL42X98d9U0THbRD3MKlB38"  # Votre clé API Gemini
-    status_code, response_json = generate_text_with_gemini(gemini_api_key, prompt)
-    email_content = response_json.get('candidates', [{}])[0].get('content', {}).get('parts', [{}])[0].get('text', 'Content generation failed.')
-
-    return draw_image, email_content
+    return draw_image, persons
 
 # Interface Streamlit
 st.title("Analyse d'Images pour Détection de Personnes")
 
+# Saisie du nombre de personnes attendues pour chaque image
+number_of_people = st.number_input("Nombre attendu de personnes pour chaque image", min_value=1, value=10, step=1)
+
 uploaded_files = st.file_uploader("Choisissez jusqu'à 3 images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
 
+# Option pour définir le seuil de confiance
+confidence_threshold = st.slider("Seuil de confiance pour la détection", 0.0, 1.0, 0.7)
+
 if uploaded_files:
+    max_absents = -1
+    best_image = None
+    best_image_name = ""
+
     for uploaded_file in uploaded_files:
         # Lire l'image téléchargée
         image = Image.open(uploaded_file)
         st.image(image, caption=f'Image téléchargée : {uploaded_file.name}', use_column_width=True)
         
         # Traitement de l'image
-        modified_image, email_content = process_image(image)
+        modified_image, detected_persons = process_image(image, confidence_threshold)
         
+        # Calculer le nombre d'absents
+        absents = number_of_people - detected_persons
+        if absents < 0:
+            absents = 0
+
         # Afficher l'image modifiée
         st.image(modified_image, caption=f'Image avec personnes détectées : {uploaded_file.name}', use_column_width=True)
         
-        # Afficher le contenu généré
-        st.write(f"**Contenu généré pour {uploaded_file.name} :**")
-        st.markdown(email_content)
+        # Afficher le nombre d'absents
+        st.write(f"**Nombre d'absents pour {uploaded_file.name} : {absents}**")
+
+        # Garder une trace de l'image avec le plus grand nombre d'absents
+        if absents > max_absents:
+            max_absents = absents
+            best_image = modified_image
+            best_image_name = uploaded_file.name
+
+    # Afficher l'image avec le plus grand nombre d'absents
+    if best_image:
+        st.image(best_image, caption=f'Image avec le plus grand nombre d\'absents : {best_image_name}', use_column_width=True)