""" tracker.py — Core detection + tracking engine Counting : FIRST-SEEN (chaque track_id compté une seule fois dès sa 1re apparition) CSV log : suit exactement le schéma SCHEMA_EXAMPLE.csv frame, timestamp_sec, scene_name, group_id, video_name, track_id, class_name, confidence, bbox_x1, bbox_y1, bbox_x2, bbox_y2, cx, cy, frame_width, frame_height, crossed_line, direction, speed_px_s """ import cv2 import csv import json import math import re import uuid import numpy as np from pathlib import Path from datetime import datetime from collections import defaultdict from ultralytics import YOLO # ─── Classes COCO traffic ───────────────────────────────────────────────────── TRAFFIC_CLASSES = { 0: "person", 1: "bicycle", 2: "car", 3: "motorbike", 5: "bus", 7: "truck", } CLASS_COLORS = { "person": (0, 200, 255), "bicycle": (50, 255, 50), "car": (255, 165, 0), "motorbike": (255, 50, 200), "bus": (0, 100, 255), "truck": (180, 0, 255), } DEFAULT_CLASSES = list(TRAFFIC_CLASSES.values()) # Colonnes CSV — ordre exact du schéma CSV_COLUMNS = [ "frame", "timestamp_sec", "scene_name", "group_id", "video_name", "track_id", "class_name", "confidence", "bbox_x1", "bbox_y1", "bbox_x2", "bbox_y2", "cx", "cy", "frame_width", "frame_height", "crossed_line", "direction", "speed_px_s", ] class TrafficTracker: def __init__( self, model_path: str = "best.pt", selected_classes: list = None, conf_threshold: float = 0.5, scene_name: str = "scene_01", group_id: str = "group_05", video_name: str = "video.mp4", output_dir: str = "logs", ): # ── Modèle ──────────────────────────────────────────────────────────── # Pour utiliser un modèle fine-tuné : # model_path = "runs/detect/traffic_finetune/weights/best.pt" self.model = YOLO(model_path) self.selected_classes = selected_classes or DEFAULT_CLASSES self.conf = conf_threshold self.scene_name = scene_name self.group_id = group_id # <— nouveau champ schéma self.video_name = video_name # <— nouveau champ schéma self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # State self.session_id = str(uuid.uuid4())[:8] self.frame_index = 0 self.fps = 30.0 self.frame_width = 0 self.frame_height = 0 # Tracking self.track_history: dict = defaultdict(list) # id -> [(cx,cy), ...] self.counted_ids: set = set() self.count_per_class: dict = defaultdict(int) # Logs self.detection_log: list = [] # chaque ligne = une détection self.frame_stats: list = [] # résumé par frame # Video writer self.video_writer = None self.output_video_path = None # Filtre classes COCO self._class_ids = [ cid for cid, name in TRAFFIC_CLASSES.items() if name in self.selected_classes ] # ── Get next order number ─────────────────────────────────────────────────── def _get_next_order_number(self) -> int: """ Retourne le prochain numéro d'ordre pour les fichiers de logs. Cherche les fichiers existants avec le pattern Group_X_Y_NNN_*. """ try: prefix = f"{self.group_id}_{self.scene_name}" pattern = f"{prefix}_*" order_pattern = re.compile(rf"^{re.escape(prefix)}_(\d{{3}})(?:_|\.|$)") existing = list(self.output_dir.glob(pattern)) if not existing: return 1 numbers = [] for f in existing: match = order_pattern.match(f.name) if match: numbers.append(int(match.group(1))) return max(numbers, default=0) + 1 except Exception: return 1 def setup_frame_source(self, width: int, height: int, fps: float = 5.0): self.fps = fps or 5.0 self.frame_width = width self.frame_height = height # ── Setup ────────────────────────────────────────────────────────────────── def setup_video(self, cap: cv2.VideoCapture, save_output: bool = True): self.fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 self.frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) if save_output: # Générer le numéro d'ordre (incrément des sessions) video_order = self._get_next_order_number() out_name = f"{self.group_id}_{self.scene_name}_{video_order:03d}_annotated.mp4" self.output_video_path = str(self.output_dir / out_name) fourcc = cv2.VideoWriter_fourcc(*"mp4v") self.video_writer = cv2.VideoWriter( self.output_video_path, fourcc, self.fps, (self.frame_width, self.frame_height) ) # ── Process one frame ────────────────────────────────────────────────────── def process_frame(self, frame: np.ndarray) -> tuple: self.frame_index += 1 timestamp_sec = round(self.frame_index / self.fps, 3) results = self.model.track( frame, persist=True, conf=self.conf, classes=self._class_ids, tracker="bytetrack.yaml", verbose=False, ) annotated = frame.copy() frame_detections = [] per_class_in_frame = defaultdict(int) any_object_visible = False result = results[0] if result.boxes is not None and len(result.boxes) > 0: for box in result.boxes: cls_id = int(box.cls[0]) cls_name = TRAFFIC_CLASSES.get(cls_id, "unknown") if cls_name not in self.selected_classes: continue conf_val = float(box.conf[0]) x1, y1, x2, y2 = map(int, box.xyxy[0]) track_id = int(box.id[0]) if box.id is not None else -1 cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 any_object_visible = True per_class_in_frame[cls_name] += 1 # ── Vitesse (px/s) ──────────────────────────────────────────── # Distance euclidienne entre position courante et précédente, # multipliée par fps pour obtenir des px/s. speed_px_s = 0.0 if track_id != -1 and self.track_history[track_id]: prev_cx, prev_cy = self.track_history[track_id][-1] dist = math.hypot(cx - prev_cx, cy - prev_cy) speed_px_s = round(dist * self.fps, 1) # ── Comptage first-seen ─────────────────────────────────────── is_new = False if track_id != -1 and track_id not in self.counted_ids: self.counted_ids.add(track_id) self.count_per_class[cls_name] += 1 is_new = True # Mise à jour historique if track_id != -1: self.track_history[track_id].append((cx, cy)) # ── Colonnes schéma : crossed_line & direction ───────────────── # On ne dessine plus de ligne physique sur la vidéo, mais on # conserve les colonnes dans le CSV (vide / false par défaut). # Si tu veux réactiver le croisement de ligne, tu peux définir # une ligne ici et compléter la logique. crossed_line = False direction = "" # ── Ligne CSV (schéma exact) ─────────────────────────────────── row = { "frame": self.frame_index, "timestamp_sec": timestamp_sec, "scene_name": self.scene_name, "group_id": self.group_id, "video_name": self.video_name, "track_id": track_id, "class_name": cls_name, "confidence": round(conf_val, 3), "bbox_x1": x1, "bbox_y1": y1, "bbox_x2": x2, "bbox_y2": y2, "cx": cx, "cy": cy, "frame_width": self.frame_width, "frame_height": self.frame_height, "crossed_line": str(crossed_line).lower(), # "true"/"false" "direction": direction, "speed_px_s": speed_px_s, } frame_detections.append(row) self.detection_log.append(row) # ── Bounding box ─────────────────────────────────────────────── color = CLASS_COLORS.get(cls_name, (255, 255, 255)) thickness = 3 if is_new else 2 cv2.rectangle(annotated, (x1, y1), (x2, y2), color, thickness) id_str = f"#{track_id}" if track_id != -1 else "#?" label = f"{cls_name} {id_str} {conf_val:.2f}" (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1) cv2.rectangle(annotated, (x1, y1 - th - 8), (x1 + tw + 4, y1), color, -1) cv2.putText(annotated, label, (x1 + 2, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 1) # Trail if track_id != -1: trail = self.track_history[track_id][-20:] for i in range(1, len(trail)): alpha = i / len(trail) tc = tuple(int(c * alpha) for c in color) cv2.line(annotated, trail[i - 1], trail[i], tc, 2) self._draw_counters(annotated, any_object_visible) if self.video_writer is not None: self.video_writer.write(annotated) frame_stat = { "frame": self.frame_index, "timestamp": timestamp_sec, "scene": self.scene_name, "detections": len(frame_detections), "per_class": dict(per_class_in_frame), "any_visible": any_object_visible, "cumulative": dict(self.count_per_class), } self.frame_stats.append(frame_stat) return annotated, frame_stat # ── Overlay compteurs ────────────────────────────────────────────────────── def _draw_counters(self, frame: np.ndarray, any_visible: bool): overlay = frame.copy() cv2.rectangle(overlay, (0, 0), (260, 30 + 22 * len(self.selected_classes) + 30), (20, 20, 20), -1) cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame) cv2.putText(frame, "TRAFFIC COUNTER", (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2) y = 42 for cls_name in self.selected_classes: count = self.count_per_class.get(cls_name, 0) color = CLASS_COLORS.get(cls_name, (255, 255, 255)) cv2.putText(frame, f" {cls_name:<12} {count:>4}", (8, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1) y += 22 if not any_visible: cv2.putText(frame, "NO OBJECTS DETECTED", (10, y + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (80, 80, 80), 1) # ── Sauvegarde des logs ──────────────────────────────────────────────────── def save_logs(self) -> dict: if self.video_writer is not None: self.video_writer.release() self.video_writer = None # Générer le numéro d'ordre pour les logs log_order = self._get_next_order_number() prefix = f"{self.group_id}_{self.scene_name}_{log_order:03d}" # ── CSV principal — schéma exact ────────────────────────────────────── csv_path = self.output_dir / f"{prefix}_detections.csv" with open(csv_path, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS) writer.writeheader() for row in self.detection_log: # S'assurer que toutes les colonnes sont présentes writer.writerow({col: row.get(col, "") for col in CSV_COLUMNS}) # ── JSONL brut (optionnel, pour debug) ──────────────────────────────── jsonl_path = self.output_dir / f"{prefix}_detections.jsonl" with open(jsonl_path, "w", encoding="utf-8") as f: for row in self.detection_log: f.write(json.dumps(row) + "\n") # ── Summary JSON ────────────────────────────────────────────────────── summary = self.get_summary() sum_path = self.output_dir / f"{prefix}_summary.json" with open(sum_path, "w", encoding="utf-8") as f: json.dump(summary, f, indent=2) # ── Frame stats CSV (résumé par frame) ──────────────────────────────── fstats_path = self.output_dir / f"{prefix}_frame_stats.csv" if self.frame_stats: with open(fstats_path, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter( f, fieldnames=self.frame_stats[0].keys() ) writer.writeheader() writer.writerows(self.frame_stats) return { "detections_csv": str(csv_path), "detections_jsonl": str(jsonl_path), "summary": str(sum_path), "frame_stats": str(fstats_path), "annotated_video": self.output_video_path or "", } # ── Summary ───────────────────────────────────────────────────────────────── def get_summary(self) -> dict: total_duration = self.frame_index / self.fps total_objects = sum(self.count_per_class.values()) buckets: dict = defaultdict(int) for stat in self.frame_stats: bucket = int(stat["timestamp"] // 10) buckets[bucket] += stat["detections"] temporal = [{"bucket_10s": k, "detections": v} for k, v in sorted(buckets.items())] return { "scene": self.scene_name, "group_id": self.group_id, "video_name": self.video_name, "session_id": self.session_id, "processed_at": datetime.now().isoformat(), "total_frames": self.frame_index, "duration_sec": round(total_duration, 2), "fps": round(self.fps, 2), "resolution": [self.frame_width, self.frame_height], "selected_classes": self.selected_classes, "total_unique_objects": total_objects, "count_per_class": dict(self.count_per_class), "annotated_video": self.output_video_path or "", "temporal_distribution": temporal, }