Update app.py

app.py

@@ -1,102 +1,218 @@
-import matplotlib
-matplotlib.use('Agg')
-
 import streamlit as st
 import cv2
-import numpy as np
-from yolov5 import YOLOv5
-from sort.sort import Sort
 import tempfile
-import shutil
-from moviepy.editor import VideoFileClip, concatenate_videoclips, ImageSequenceClip
 import os
-
-# Load the pre-trained model and initialize the SORT tracker
-model_path = 'yolov5s.pt'  # Ensure this path points to the model file
-model = YOLOv5(model_path, device='cpu')
-tracker = Sort()
-
-def process_video(uploaded_file):
-    # Save the uploaded file to a temporary location
-    temp_file_path = "temp_video.mp4"
-    with open(temp_file_path, "wb") as temp_file:
-        temp_file.write(uploaded_file.getbuffer())
-
-    # Use moviepy to read the video file
-    video_clip = VideoFileClip(temp_file_path)
-    total_frames = int(video_clip.fps * video_clip.duration)
-    width, height = video_clip.size
-
-    # Temporary directory to save processed video frames
-    temp_dir = tempfile.mkdtemp()
-
-    unique_cars = set()
-    progress_bar = st.progress(0)
-
-    for frame_idx, frame in enumerate(video_clip.iter_frames()):
-        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-
-        progress_percentage = int((frame_idx + 1) / total_frames * 100)
-        progress_bar.progress(progress_percentage)
-
-        # Detection and tracking logic
-        results = model.predict(frame)
-        preds = results.pandas().xyxy[0]
-        detections = []
-
-        for index, row in preds.iterrows():
-            if row['name'] == 'car':
-                xmin, ymin, xmax, ymax, conf = row['xmin'], row['ymin'], row['xmax'], row['ymax'], row['confidence']
-                detections.append([xmin, ymin, xmax, ymax, conf])
-
-        if detections:
-            detections_np = np.array(detections)
-            trackers = tracker.update(detections_np)
-
-            for d in trackers:
-                unique_cars.add(int(d[4]))
-                xmin, ymin, xmax, ymax = map(int, d[:4])
-                cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
-                cv2.putText(frame, f'ID: {int(d[4])}', (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
-
-        cv2.putText(frame, f'Unique Cars: {len(unique_cars)}', (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0, 255, 0), 2)
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert back to RGB for moviepy
-        cv2.imwrite(f"{temp_dir}/{frame_idx:04d}.jpg", frame)
-
-    frames_files = [os.path.join(temp_dir, f"{i:04d}.jpg") for i in range(total_frames)]
-    clip = ImageSequenceClip(frames_files, fps=video_clip.fps)
-    output_video_path = 'processed_video.mp4'
-    clip.write_videofile(output_video_path, codec='libx264')  # Use libx264 codec for compatibility
-
-    # Remove temporary directory and temporary files
-    shutil.rmtree(temp_dir)
-
-    return output_video_path
-
-
+import time
+import numpy as np
+import pandas as pd
+from collections import defaultdict
+from ultralytics import YOLO
+
+# --- FONCTIONS UTILES ---
+def draw_text_with_background(image, text, position, font=cv2.FONT_HERSHEY_SIMPLEX,
+                              font_scale=1, font_thickness=2, text_color=(255, 255, 255), bg_color=(0, 0, 0), padding=5):
+    """Ajoute du texte avec un fond sur une image OpenCV."""
+    text_size = cv2.getTextSize(text, font, font_scale, font_thickness)[0]
+    text_width, text_height = text_size
+
+    x, y = position
+    top_left = (x, y - text_height - padding)
+    bottom_right = (x + text_width + padding * 2, y + padding)
+
+    cv2.rectangle(image, top_left, bottom_right, bg_color, -1)
+    cv2.putText(image, text, (x + padding, y), font, font_scale, text_color, font_thickness, cv2.LINE_AA)
+
+# --- CLASSE YOLO ---
+class YOLOVideoProcessor:
+    def __init__(self, model_path, video_path, output_path, poly1, poly2, tracker_method="bot"):
+        self.model = YOLO(model_path, task="detect")
+        self.tracker_method = tracker_method
+        self.video_path = video_path
+        self.output_path = output_path 
+
+        self.unique_region1_ids = set()
+        self.unique_region2_ids = set()
+        self.poly1 = poly1
+        self.poly2 = poly2
+
+    def is_in_region(self, center, poly):
+        poly_np = np.array(poly, dtype=np.int32)
+        return cv2.pointPolygonTest(poly_np, center, False) >= 0
+
+    def process_video(self, progress_bar=None):
+        cap = cv2.VideoCapture(self.video_path)
+        if not cap.isOpened():
+            st.error("⚠️ Erreur : Impossible d'ouvrir la vidéo.")
+            return
+
+        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+        if fps == 0:
+            fps = 30  # Valeur par défaut si FPS est invalide
+
+        # Utiliser XVID qui est généralement mieux supporté
+        fourcc = cv2.VideoWriter_fourcc(*'XVID') 
+        out = cv2.VideoWriter(self.output_path, fourcc, fps, (frame_width, frame_height))  
+
+        processed_frames = 0
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        while cap.isOpened():
+            success, frame = cap.read()
+            if not success:
+                break
+
+            # Mise à jour de la barre de progression
+            if progress_bar is not None:
+                progress_bar.progress(processed_frames / total_frames)
+
+            tracker = "botsort.yaml" if self.tracker_method.lower() == "bot" else "bytetrack.yaml"
+            results = self.model.track(frame, persist=True, tracker=tracker, conf=0.25)
+
+            track_ids = []
+            if results and len(results) > 0 and len(results[0].boxes) > 0:
+                try:
+                    track_ids = results[0].boxes.id.int().cpu().tolist()
+                except AttributeError:
+                    track_ids = [i for i in range(len(results[0].boxes.xywh.cpu().numpy()))]
+
+            # Dessiner les polygones
+            cv2.polylines(frame, [np.array(self.poly1, np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)
+            cv2.polylines(frame, [np.array(self.poly2, np.int32)], isClosed=True, color=(255, 0, 0), thickness=2)
+
+            for box, track_id in zip(results[0].boxes.xywh.cpu().numpy(), track_ids):
+                x, y, w, h = box
+                center_point = (int(x), int(y))
+
+                if self.is_in_region(center_point, self.poly1):
+                    self.unique_region1_ids.add(track_id)
+                if self.is_in_region(center_point, self.poly2):
+                    self.unique_region2_ids.add(track_id)
+
+            # Affichage du comptage des véhicules
+            draw_text_with_background(frame, f'Total Sens 1: {len(self.unique_region1_ids)}', (10, frame_height - 50))
+            draw_text_with_background(frame, f'Total Sens 2: {len(self.unique_region2_ids)}', (frame_width - 300, frame_height - 50))
+
+            out.write(frame)
+            processed_frames += 1
+
+        cap.release()
+        out.release()
+        cv2.destroyAllWindows()
+
+        if processed_frames == 0:
+            st.error("⚠️ Aucune image n'a été écrite dans la vidéo de sortie !")
+        
+        return len(self.unique_region1_ids), len(self.unique_region2_ids)
+
+
+# --- INTERFACE STREAMLIT ---
 def main():
-    # Initialize session state variables if they don't exist
-    if 'output_video_path' not in st.session_state:
-        st.session_state.output_video_path = None
-
-    st.sidebar.image("logo.jpg", use_column_width=True)
-    uploaded_file = st.sidebar.file_uploader("Upload a video", type=['mp4'])
-
-    st.title("Car Detection and Tracking")
-
-    if uploaded_file is not None:
-        # Process the video only if it hasn't been processed yet or a new file is uploaded
-        if st.session_state.output_video_path is None or st.session_state.uploaded_file_name != uploaded_file.name:
-            st.session_state.uploaded_file_name = uploaded_file.name
-            st.session_state.output_video_path = process_video(uploaded_file)
-
-        # Display the processed video
-        st.video(st.session_state.output_video_path)
-
-        # Provide a download link for the processed video
-        with open(st.session_state.output_video_path, "rb") as file:
-            st.download_button("Download Processed Video", file, file_name="processed_video.mp4")
-
+    st.set_page_config(
+        page_title="Détecteur de Véhicules",
+        page_icon="🚗",
+        layout="wide"
+    )
+    
+    st.title("🚗 Détection et comptage de Véhicules sur l'Autoroute de l'Avenir")
+    
+    # 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")
+                st.success("✅ Modèle chargé avec succès!")
+            except Exception as e:
+                st.error(f"❌ Erreur lors du chargement du modèle: {e}")
+                # Fallback: utiliser un modèle YOLO standard
+                st.warning("⚠️ Utilisation du modèle YOLO standard à la place")
+                model_path = "yolov8n.pt"
+    
+    # Colonnes pour l'organisation de l'interface
+    col1, col2 = st.columns([3, 1])
+    
+    with col2:
+        st.header("🔹 Paramètres")
+        
+        # Entrée utilisateur pour les polygones
+        st.subheader("📍 Polygone 1 (vert)")
+        poly1_input = st.text_area("Entrez 4 points (x,y) séparés par des espaces", "465,350 609,350 520,630 3,630")
+        
+        st.subheader("📍 Polygone 2 (rouge)")
+        poly2_input = st.text_area("Entrez 4 points (x,y) séparés par des espaces", "678,350 815,350 1203,630 743,630")
+        
+        tracker_method = st.selectbox("Méthode de tracking", ["bot", "byte"], index=0)
+    
+    with col1:
+        uploaded_file = st.file_uploader("📂 Upload une vidéo", type=["mp4", "avi", "mov"])
+        
+        def parse_polygon(input_text):
+            try:
+                return [tuple(map(int, point.split(','))) for point in input_text.split()]
+            except:
+                return []
+        
+        poly1 = parse_polygon(poly1_input)
+        poly2 = parse_polygon(poly2_input)
+        
+        if uploaded_file is not None:
+            # Créer un dossier temporaire si nécessaire
+            temp_dir = tempfile.mkdtemp()
+            input_video_path = os.path.join(temp_dir, "input_video.mp4")
+            output_video_path = os.path.join(temp_dir, "output_video.mp4")
+            
+            # Écrire le fichier téléchargé dans un fichier temporaire
+            with open(input_video_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+            
+            st.video(input_video_path)  # Afficher la vidéo d'entrée
+            
+            if st.button("▶️ Lancer la détection"):
+                if len(poly1) == 4 and len(poly2) == 4:
+                    # Afficher la barre de progression
+                    progress_text = "🔄 Traitement de la vidéo en cours..."
+                    progress_bar = st.progress(0)
+                    
+                    # Traitement de la vidéo
+                    processor = YOLOVideoProcessor(model_path, input_video_path, output_video_path, poly1, poly2, tracker_method)
+                    
+                    # Démarrer le traitement
+                    start_time = time.time()
+                    count1, count2 = processor.process_video(progress_bar=progress_bar)
+                    end_time = time.time()
+                    
+                    # Calcul du temps de traitement
+                    processing_time = end_time - start_time
+                    
+                    progress_bar.progress(1.0)  # Compléter la barre de progression
+                    st.success(f"✅ Traitement terminé en {processing_time:.2f} secondes!")
+                    
+                    # Afficher les résultats
+                    col_result1, col_result2 = st.columns(2)
+                    with col_result1:
+                        st.metric("Véhicules Sens 1 (Vert)", count1)
+                    with col_result2:
+                        st.metric("Véhicules Sens 2 (Rouge)", count2)
+                    
+                    # Afficher la vidéo traitée
+                    st.subheader("Vidéo traitée")
+                    st.video(output_video_path)
+                    
+                    # Option de téléchargement
+                    with open(output_video_path, "rb") as file:
+                        st.download_button(
+                            label="⬇️ Télécharger la vidéo",
+                            data=file,
+                            file_name="video_traitee.mp4",
+                            mime="video/mp4"
+                        )
+                else:
+                    st.error("❌ Les coordonnées des polygones doivent contenir **exactement 4 points**.")
 
 if __name__ == "__main__":
-    main()
+    main()