Update app.py

app.py

@@ -248,54 +248,67 @@ class YOLOVideoProcessor:
             
             frame_count += 1
             
-            # Attendre un peu moins pour augmenter le FPS
-            time.sleep(0.01)
-        
         cap.release()
         st.success("✅ Flux vidéo arrêté.")
 
 
-# --- NOUVELLE FONCTION POUR AFFICHER PRÉVISUALISATION DE LA WEBCAM ---
+# --- FONCTION CORRIGÉE POUR AFFICHER PRÉVISUALISATION DE LA WEBCAM ---
 def display_webcam_preview(camera_id, display_placeholder, poly1=None, poly2=None):
     """Affiche une prévisualisation simple de la webcam sans détection"""
+    # Utiliser directement OpenCV au lieu d'un thread
     cap = cv2.VideoCapture(camera_id)
     
     # Configuration de la webcam pour une bonne performance
     cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
     cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
-    cap.set(cv2.CAP_PROP_FPS, 30)
     
     if not cap.isOpened():
         st.error("⚠️ Erreur : Impossible d'ouvrir la webcam.")
         return False
     
-    stop_preview = False
+    # Lecture d'une seule frame pour vérifier que tout fonctionne
+    success, frame = cap.read()
+    if not success:
+        st.error("⚠️ Erreur lors de la lecture du flux vidéo.")
+        cap.release()
+        return False
+    
+    # Initialiser st.session_state pour le contrôle de la prévisualisation
+    if 'preview_stop' not in st.session_state:
+        st.session_state.preview_stop = False
+    
+    # Boucle principale de la prévisualisation
     st.session_state.preview_active = True
     
-    while st.session_state.preview_active and not stop_preview:
-        success, frame = cap.read()
-        if not success:
-            st.error("⚠️ Erreur lors de la lecture du flux vidéo.")
-            break
-        
-        # Dessiner les polygones si disponibles
-        if poly1 and len(poly1) >= 3:
-            cv2.polylines(frame, [np.array(poly1, np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
-        if poly2 and len(poly2) >= 3:
-            cv2.polylines(frame, [np.array(poly2, np.int32)], isClosed=True, color=(255, 0, 0), thickness=2)
-        
-        # Ajouter un texte explicatif
-        draw_text_with_background(frame, "Prévisualisation (sans détection)", (10, 30))
-        
-        # Convertir et afficher l'image
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        img = Image.fromarray(frame_rgb)
-        display_placeholder.image(img, channels="RGB", use_column_width=True)
-        
-        # Petit délai pour ne pas surcharger l'interface
-        time.sleep(0.03)
+    try:
+        while st.session_state.preview_active and not st.session_state.preview_stop:
+            success, frame = cap.read()
+            if not success:
+                break
+            
+            # Dessiner les polygones si disponibles
+            if poly1 and len(poly1) >= 3:
+                cv2.polylines(frame, [np.array(poly1, np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
+            if poly2 and len(poly2) >= 3:
+                cv2.polylines(frame, [np.array(poly2, np.int32)], isClosed=True, color=(255, 0, 0), thickness=2)
+            
+            # Ajouter un texte explicatif
+            draw_text_with_background(frame, "Prévisualisation (sans détection)", (10, 30))
+            
+            # Convertir et afficher l'image
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            img = Image.fromarray(frame_rgb)
+            
+            # Afficher l'image dans le placeholder
+            display_placeholder.image(img, channels="RGB", use_column_width=True)
+            
+            # Court délai pour ne pas saturer Streamlit
+            time.sleep(0.1)
+    except Exception as e:
+        st.error(f"Erreur lors de la prévisualisation: {e}")
+    finally:
+        cap.release()
     
-    cap.release()
     return True
 
 
@@ -309,21 +322,20 @@ def main():
     
     st.title("🚗 Détection et comptage de Véhicules sur l'Autoroute de l'Avenir")
     
-    # Session state pour gérer l'état de la webcam et de la prévisualisation
+    # Initialisation des variables d'état de session
     if 'webcam_active' not in st.session_state:
         st.session_state.webcam_active = False
     if 'preview_active' not in st.session_state:
         st.session_state.preview_active = False
+    if 'preview_stop' not in st.session_state:
+        st.session_state.preview_stop = False
     if 'processor' not in st.session_state:
         st.session_state.processor = None
-    if 'preview_thread' not in st.session_state:
-        st.session_state.preview_thread = None
     
     # Vérifier si le modèle existe déjà ou doit être téléchargé
     model_path = "best.pt"
     if not os.path.exists(model_path):
         with st.spinner("📥 Chargement du modèle YOLO... Cela peut prendre un moment."):
-            # Utilisez hub.load pour télécharger le modèle depuis Hugging Face Hub
             try:
                 from huggingface_hub import hf_hub_download
                 model_path = hf_hub_download(repo_id="ModuMLTECH/projet_trafic_2", filename="best.pt")
@@ -436,7 +448,7 @@ def main():
         # Sélectionner la source vidéo
         camera_options = {"Webcam par défaut": 0}
         # Ajouter des options pour d'autres caméras si disponibles
-        for i in range(1, 5):  # Essayer de détecter jusqu'à 5 caméras
+        for i in range(1, 3):  # Limiter à 3 caméras pour ne pas ralentir le chargement
             try:
                 cap = cv2.VideoCapture(i)
                 if cap.isOpened():
@@ -458,29 +470,36 @@ def main():
         preview_col1, preview_col2 = st.columns(2)
         
         with preview_col1:
-            if st.button("👁️ Afficher la prévisualisation"):
+            if st.button("👁️ Afficher la prévisualisation", key="start_preview"):
                 # Arrêter la détection si elle est active
-                if st.session_state.webcam_active:
+                if st.session_state.webcam_active and hasattr(st.session_state, 'processor'):
                     st.session_state.processor.stop_processing = True
                     st.session_state.webcam_active = False
                     time.sleep(0.5)  # Attendre que le thread se termine
                 
-                # Démarrer la prévisualisation
-                if not st.session_state.preview_active:
-                    st.session_state.preview_active = True
-                    # Démarrer la prévisualisation dans un thread séparé
-                    st.session_state.preview_thread = threading.Thread(
-                        target=display_webcam_preview,
-                        args=(camera_id, video_placeholder, poly1, poly2)
-                    )
-                    st.session_state.preview_thread.daemon = True
-                    st.session_state.preview_thread.start()
+                # Réinitialiser le flag d'arrêt
+                st.session_state.preview_stop = False
+                st.session_state.preview_active = True
+                
+                # Utiliser un thread pour la prévisualisation
+                preview_thread = threading.Thread(
+                    target=display_webcam_preview,
+                    args=(camera_id, video_placeholder, poly1, poly2)
+                )
+                preview_thread.daemon = True
+                preview_thread.start()
+                
+                # Message de succès
+                st.success("📷 Prévisualisation démarrée")
                     
         with preview_col2:
-            if st.button("⏹️ Arrêter la prévisualisation"):
+            if st.button("⏹️ Arrêter la prévisualisation", key="stop_preview"):
+                # Arrêter la prévisualisation
+                st.session_state.preview_stop = True
                 st.session_state.preview_active = False
                 time.sleep(0.5)  # Attendre que le thread se termine
                 video_placeholder.empty()  # Effacer l'affichage vidéo
+                st.success("🛑 Prévisualisation arrêtée")
                 
         # Section de détection
         st.subheader("2️⃣ Détection des véhicules")
@@ -492,7 +511,7 @@ def main():
         detect_col1, detect_col2 = st.columns(2)
         
         with detect_col1:
-            if st.button("▶️ Démarrer la détection"):
+            if st.button("▶️ Démarrer la détection", key="start_detection"):
                 if not valid_polygons:
                     st.error("❌ Les coordonnées des polygones doivent contenir **exactement 4 points**.")
                 elif st.session_state.webcam_active:
@@ -500,6 +519,7 @@ def main():
                 else:
                     # Arrêter la prévisualisation si elle est active
                     if st.session_state.preview_active:
+                        st.session_state.preview_stop = True
                         st.session_state.preview_active = False
                         time.sleep(0.5)  # Attendre que le thread se termine
                     
@@ -514,19 +534,22 @@ def main():
                     
                     # Démarrer le traitement dans un thread séparé
                     processing_thread = threading.Thread(
-                        target=st.session_state.processor.process_webcam,
+                        target=processor.process_webcam,
                         args=(camera_id, video_placeholder, count_placeholders)
                     )
                     processing_thread.daemon = True
                     processing_thread.start()
+                    
+                    st.success("✅ Détection démarrée")
                 
         with detect_col2:
-            if st.button("⏹️ Arrêter la détection"):
-                if st.session_state.webcam_active and st.session_state.processor:
+            if st.button("⏹️ Arrêter la détection", key="stop_detection"):
+                if st.session_state.webcam_active and hasattr(st.session_state, 'processor'):
                     st.session_state.processor.stop_processing = True
                     st.session_state.webcam_active = False
                     time.sleep(0.5)  # Attendre que le thread se termine
                     video_placeholder.empty()  # Effacer l'affichage vidéo
+                    st.success("🛑 Détection arrêtée")
                 else:
                     st.warning("⚠️ Aucune détection en cours !")