Upload 3 files

utils/ssd_detector.py

@@ -0,0 +1,190 @@
+"""
+SSD MobileNetV3 – Object detection via torchvision.
+Detects traffic objects AND stop signs.
+GPU forced when available.
+No external model files needed – weights auto-downloaded.
+"""
+
+import cv2 as cv
+import torch
+import torchvision
+from torchvision.transforms import functional as F
+import numpy as np
+import os
+import csv
+from datetime import datetime
+
+COCO_CLASSES = [
+    "__background__", "person", "bicycle", "car", "motorcycle", "airplane",
+    "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
+    "N/A", "stop sign", "parking meter", "bench", "bird", "cat", "dog",
+    "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "N/A",
+    "backpack", "umbrella", "N/A", "N/A", "handbag", "tie", "suitcase",
+    "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
+    "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
+    "N/A", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana",
+    "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza",
+    "donut", "cake", "chair", "couch", "potted plant", "bed", "N/A",
+    "dining table", "N/A", "N/A", "toilet", "N/A", "tv", "laptop", "mouse",
+    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
+    "sink", "refrigerator", "N/A", "book", "clock", "vase", "scissors",
+    "teddy bear", "hair drier", "toothbrush"
+]
+
+# ✅ "motorbike" aligné avec yolo_tracker.py (cohérence du projet)
+TRAFFIC_CLASSES = {
+    "person", "bicycle", "car", "motorbike", "motorcycle",
+    "bus", "truck", "stop sign", "traffic light"
+}
+
+np.random.seed(42)
+COLORS     = np.random.randint(0, 255, size=(len(COCO_CLASSES), 3), dtype="uint8")
+STOP_COLOR = (0, 0, 255)
+
+
+class SSDDetector:
+    """
+    SSDLite MobileNetV3 pre-trained on COCO.
+    - No training needed
+    - Detects 91 COCO classes including stop sign
+    - GPU accelerated
+    """
+    def __init__(self, filepath,
+                 classes=None,
+                 confidence_threshold=0.4,
+                 device="cuda",
+                 output_dir="logs"):
+
+        self.filepath   = filepath
+        self.classes    = classes if classes else TRAFFIC_CLASSES
+        self.conf_thr   = confidence_threshold
+        self.output_dir = output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs("results",  exist_ok=True)
+
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        print(f"Loading SSD MobileNetV3 on {self.device}...")
+
+        weights = torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights.COCO_V1
+        self.model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(weights=weights)
+        self.model.to(self.device)
+        self.model.eval()
+        print(f"Model ready on : {self.device}")
+
+    def forward(self, show=True, save_video=True):
+        cap = cv.VideoCapture(self.filepath)
+        assert cap.isOpened(), "Cannot open video"
+
+        w   = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
+        h   = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
+        fps = cap.get(cv.CAP_PROP_FPS) or 25
+
+        video_writer = None
+        if save_video:
+            video_writer = cv.VideoWriter(
+                "results/ssd_detection.avi",
+                cv.VideoWriter_fourcc(*"mp4v"),
+                fps, (w, h)
+            )
+
+        log_path = os.path.join(
+            self.output_dir,
+            f"ssd_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+        )
+
+        stats     = {}
+        frame_idx = 0
+
+        print(f"Running SSD detection...")
+        print(f"Video : {w}x{h} @ {fps:.0f}fps")
+
+        with open(log_path, "w", newline="") as f:
+            writer = csv.writer(f)
+            # ✅ Ajout timestamp_s pour cohérence avec yolo_tracker logs
+            writer.writerow([
+                "timestamp_s", "frame", "class",
+                "x1", "y1", "x2", "y2", "confidence"
+            ])
+
+            while cap.isOpened():
+                success, frame = cap.read()
+                if not success:
+                    break
+
+                timestamp = round(frame_idx / fps, 3)
+                img_rgb   = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
+                tensor    = F.to_tensor(img_rgb).unsqueeze(0).to(self.device)
+
+                with torch.no_grad():
+                    outputs = self.model(tensor)[0]
+
+                boxes  = outputs["boxes"].cpu().numpy()
+                labels = outputs["labels"].cpu().numpy()
+                scores = outputs["scores"].cpu().numpy()
+
+                no_object = True
+
+                for box, label, score in zip(boxes, labels, scores):
+                    if score < self.conf_thr:
+                        continue
+
+                    cls_name = COCO_CLASSES[label] if label < len(COCO_CLASSES) else "unknown"
+                    # ✅ Normaliser "motorcycle" → "motorbike"
+                    if cls_name == "motorcycle":
+                        cls_name = "motorbike"
+                    if cls_name not in self.classes:
+                        continue
+
+                    no_object = False
+                    x1, y1, x2, y2 = map(int, box)
+
+                    if cls_name == "stop sign":
+                        color     = STOP_COLOR
+                        thickness = 3
+                        cv.putText(frame, "STOP SIGN",
+                                   (x1, max(y1 - 20, 0)),
+                                   cv.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+                    else:
+                        color     = [int(c) for c in COLORS[label]]
+                        thickness = 2
+
+                    cv.rectangle(frame, (x1, y1), (x2, y2), color, thickness)
+                    cv.putText(frame, f"{cls_name}: {score:.2f}",
+                               (x1, max(y1 - 8, 0)),
+                               cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+                    stats[cls_name] = stats.get(cls_name, 0) + 1
+                    writer.writerow([
+                        timestamp, frame_idx, cls_name,
+                        x1, y1, x2, y2, round(float(score), 3)
+                    ])
+
+                if no_object:
+                    cv.putText(frame, "No selected object detected",
+                               (20, 50), cv.FONT_HERSHEY_SIMPLEX,
+                               1.0, (0, 0, 255), 2)
+
+                if save_video and video_writer:
+                    video_writer.write(frame)
+                if show:
+                    cv.imshow("SSD Detection", frame)
+                    if cv.waitKey(1) & 0xFF == ord("q"):
+                        break
+
+                frame_idx += 1
+
+        cap.release()
+        if video_writer:
+            video_writer.release()
+        if show:
+            try:
+                cv.destroyAllWindows()
+            except cv.error:
+                pass
+
+        print(f"\n=== SSD Results ===")
+        print(f"Log CSV : {log_path}")
+        for cls, count in sorted(stats.items()):
+            print(f"  {cls} : {count} detections")
+
+        return log_path, stats

utils/yolo_tracker.py

@@ -0,0 +1,182 @@
+"""
+YOLOv11 + ByteTrack – Detection, tracking, and CSV logging.
+Detects traffic objects AND stop signs.
+GPU forced when available.
+"""
+
+import cv2 as cv
+from ultralytics import YOLO
+import csv
+import os
+from datetime import datetime
+
+TRAFFIC_CLASSES = [
+    'car', 'truck', 'bus', 'motorbike', 'bicycle',
+    'person', 'traffic sign', 'traffic light'
+]
+
+CLASS_ALIASES = {
+    "motorcycle": "motorbike",
+    "trafficLight": "traffic light",
+    "traffic_light": "traffic light",
+    "pedestrian": "person",
+}
+
+
+def normalize_class_name(name: str) -> str:
+    return CLASS_ALIASES.get(name, name)
+
+
+class Tracker:
+    """
+    YOLOv11 + ByteTrack tracker.
+    - Assigns unique IDs to each object
+    - Logs every detection to CSV
+    - Highlights stop signs with a special color (red)
+    - GPU accelerated
+    """
+    def __init__(self, filepath, classes=None, device="cuda",
+                 output_dir="logs", conf=0.4, min_box_area=900, min_track_hits=3):
+
+        self.filepath   = filepath
+        self.classes    = classes if classes else TRAFFIC_CLASSES
+        self.device     = device
+        self.output_dir = output_dir
+        self.conf       = conf
+        self.min_box_area = min_box_area
+        self.min_track_hits = min_track_hits
+        os.makedirs(output_dir, exist_ok=True)
+
+        # ✅ Fallback : utilise best.pt si disponible, sinon yolo11n.pt
+        finetuned = "models/best.pt"
+        base      = "models/yolo11n.pt"
+        self.model_path = finetuned if os.path.exists(finetuned) else base
+        print(f"[Tracker] Model : {self.model_path}")
+
+    def forward(self, show=True, save_video=True):
+        model = YOLO(self.model_path)
+        cap   = cv.VideoCapture(self.filepath)
+        assert cap.isOpened(), "Cannot open video"
+
+        w     = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
+        h     = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
+        fps   = cap.get(cv.CAP_PROP_FPS) or 25
+        total = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
+
+        video_writer = None
+        if save_video:
+            os.makedirs("results", exist_ok=True)
+            video_writer = cv.VideoWriter(
+                "results/yolo_tracked.avi",
+                cv.VideoWriter_fourcc(*"mp4v"), fps, (w, h)
+            )
+
+        log_path = os.path.join(
+            self.output_dir,
+            f"log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+        )
+
+        unique_ids = {}
+        track_hits = {}
+        frame_idx  = 0
+
+        print(f"Running YOLOv11 + ByteTrack on {self.device.upper()}...")
+        print(f"Video : {w}x{h} @ {fps:.0f}fps ({total} frames)")
+
+        with open(log_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                'timestamp_s', 'frame', 'track_id',
+                'class', 'x1', 'y1', 'x2', 'y2', 'confidence'
+            ])
+
+            while cap.isOpened():
+                success, frame = cap.read()
+                if not success:
+                    break
+
+                timestamp = round(frame_idx / fps, 3)
+
+                results = model.track(
+                    frame,
+                    persist=True,
+                    tracker="bytetrack.yaml",
+                    conf=self.conf,       # ✅ seuil de confiance
+                    device=self.device,
+                    verbose=False
+                )
+
+                no_object = True
+
+                if results[0].boxes is not None and results[0].boxes.id is not None:
+                    for box in results[0].boxes:
+                        track_id = int(box.id[0])
+                        cls_name = normalize_class_name(model.names[int(box.cls[0])])
+                        conf_val = round(float(box.conf[0]), 3)
+                        x1, y1, x2, y2 = map(int, box.xyxy[0])
+                        box_area = max(0, (x2 - x1)) * max(0, (y2 - y1))
+
+                        if cls_name not in self.classes:
+                            continue
+                        if conf_val < self.conf:
+                            continue
+                        if box_area < self.min_box_area:
+                            continue
+
+                        no_object = False
+                        track_hits[track_id] = track_hits.get(track_id, 0) + 1
+                        if track_hits[track_id] >= self.min_track_hits:
+                            unique_ids.setdefault(track_id, cls_name)
+
+                        writer.writerow([
+                            timestamp, frame_idx, track_id,
+                            cls_name, x1, y1, x2, y2, conf_val
+                        ])
+
+                annotated = results[0].plot()
+
+                # Boîte rouge spéciale pour stop signs
+                if results[0].boxes is not None:
+                    for box in results[0].boxes:
+                        cls_name = normalize_class_name(model.names[int(box.cls[0])])
+                        if cls_name == 'stop sign':
+                            x1, y1, x2, y2 = map(int, box.xyxy[0])
+                            cv.rectangle(annotated, (x1, y1), (x2, y2), (0, 0, 255), 3)
+                            cv.putText(annotated, "STOP SIGN",
+                                       (x1, y1 - 10),
+                                       cv.FONT_HERSHEY_SIMPLEX, 0.8,
+                                       (0, 0, 255), 2)
+
+                if no_object:
+                    cv.putText(annotated, "No selected object detected",
+                               (20, 50), cv.FONT_HERSHEY_SIMPLEX,
+                               1.0, (0, 0, 255), 2)
+
+                if save_video and video_writer:
+                    video_writer.write(annotated)
+                if show:
+                    cv.imshow("YOLOv11 Tracking", annotated)
+                    if cv.waitKey(1) & 0xFF == ord('q'):
+                        break
+
+                frame_idx += 1
+
+        cap.release()
+        if video_writer:
+            video_writer.release()
+        if show:
+            try:
+                cv.destroyAllWindows()
+            except cv.error:
+                pass
+
+        stats = {}
+        for cls in unique_ids.values():
+            stats[cls] = stats.get(cls, 0) + 1
+
+        print(f"\n=== Results ===")
+        print(f"Log CSV : {log_path}")
+        for cls, count in sorted(stats.items()):
+            print(f"  {cls} : {count} unique objects")
+
+        return log_path, stats

utils/yolo_video.py

@@ -0,0 +1,79 @@
+"""
+YOLOv11 – Simple detection (no tracking).
+Counts objects per frame inside a region of interest.
+GPU forced when available.
+"""
+
+import cv2 as cv
+from ultralytics import solutions
+import os
+
+DETECTION_CLASSES = [
+    'car', 'truck', 'bus', 'motorbike', 'bicycle',
+    'person', 'traffic sign', 'traffic light'
+]
+
+
+class Detector:
+    def __init__(self, filepath, device="cuda"):
+        self.filepath = filepath
+        self.device   = device
+        os.makedirs("results", exist_ok=True)
+
+        
+        finetuned = "models/best.pt"
+        base      = "models/yolo11n.pt"
+        self.model_path = finetuned if os.path.exists(finetuned) else base
+        print(f"[Detector] Model : {self.model_path}")
+
+    def forward(self, show=True):
+        cap = cv.VideoCapture(self.filepath)
+        assert cap.isOpened(), "Cannot open video/image"
+
+        w   = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
+        h   = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
+        fps = int(cap.get(cv.CAP_PROP_FPS)) or 25
+
+        #  region_points dynamiques selon la résolution de la vidéo
+        margin_x = int(w * 0.02)
+        margin_y = int(h * 0.05)
+        top_y    = int(h * 0.35)
+        bot_y    = int(h * 0.80)
+        region_points = [
+            (margin_x, bot_y),
+            (w - margin_x, bot_y),
+            (w - margin_x, top_y),
+            (margin_x, top_y)
+        ]
+
+        video_writer = cv.VideoWriter(
+            "results/yolo_detection.avi",
+            cv.VideoWriter_fourcc(*"mp4v"),
+            fps, (w, h)
+        )
+
+        counter = solutions.ObjectCounter(
+            show=show,
+            region=region_points,
+            model=self.model_path,
+            device=self.device,
+        )
+
+        print(f"Running YOLOv11 detection on {self.device.upper()}...")
+        print(f"Video : {w}x{h} @ {fps}fps | Region : {region_points}")
+
+        while cap.isOpened():
+            success, frame = cap.read()
+            if not success:
+                break
+            results = counter(frame)
+            video_writer.write(results.plot_im)
+
+        cap.release()
+        video_writer.release()
+        if show:
+            try:
+                cv.destroyAllWindows()
+            except cv.error:
+                pass
+        print("Done → results/yolo_detection.avi")