Update

backend/main.py

@@ -308,3 +308,24 @@ async def download_log(filename: str):
     if not p.exists():
         raise HTTPException(404)
     return FileResponse(str(p), filename=filename)
+
+
+@app.get("/download/video/{sid}")
+async def download_annotated_video(sid: str):
+    """
+    Download the annotated output video (.mp4) for a completed session.
+    Only available after status == 'done'.
+    """
+    if sid not in sessions:
+        raise HTTPException(404, "Session not found")
+    s = sessions[sid]
+    if s["status"] != "done":
+        raise HTTPException(202, "Processing not complete yet")
+    video_path = (s.get("summary") or {}).get("annotated_video", "")
+    if not video_path or not Path(video_path).exists():
+        raise HTTPException(404, "Annotated video not found")
+    return FileResponse(
+        str(video_path),
+        media_type="video/mp4",
+        filename=Path(video_path).name,
+    )

backend/tracker.py

@@ -1,14 +1,16 @@
 """
 tracker.py — Core detection + tracking engine
 Uses YOLOv8 (ultralytics) with ByteTrack built-in tracker.
-Supports: cars, buses, trucks, motorbikes, bicycles, pedestrians
 
 Counting strategy: FIRST-SEEN — each unique track_id is counted exactly
-once, on the first frame it appears. No counting line needed.
-Untracked detections (track_id == -1) are drawn but never counted.
+once, on the first frame it appears. No counting line.
+
+Video output: every processed frame (with bounding boxes) is written to
+an .mp4 file in the logs/ directory via cv2.VideoWriter.
 """
 
 import cv2
+import csv
 import json
 import uuid
 import numpy as np
@@ -39,17 +41,22 @@ CLASS_COLORS = {
 DEFAULT_CLASSES = list(TRAFFIC_CLASSES.values())
 
 
-# ─── TrafficTracker ────────────────────────────────────────────────────────────
 class TrafficTracker:
     def __init__(
         self,
         model_path: str = "yolov8n.pt",
-        selected_classes: list[str] = None,
+        selected_classes: list = None,
         conf_threshold: float = 0.35,
         scene_name: str = "scene_01",
         output_dir: str = "logs",
     ):
+        # ── Model ─────────────────────────────────────────────────────────────
+        # model_path is selected in the frontend dropdown (main.py receives it
+        # as a form field and passes it here).
+        # To use a fine-tuned model, pass the path to your best.pt here, e.g.:
+        #   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
@@ -63,16 +70,18 @@ class TrafficTracker:
         self.frame_width  = 0
         self.frame_height = 0
 
-        # Tracking state
-        # track_history: id -> list of (cx, cy) centres (for trail drawing)
-        self.track_history:   dict[int, list]  = defaultdict(list)
-        # counted_ids: set of track_ids already counted (never count twice)
-        self.counted_ids:     set[int]         = set()
-        self.count_per_class: dict[str, int]   = defaultdict(int)
+        # Tracking
+        self.track_history:   dict = defaultdict(list)
+        self.counted_ids:     set  = set()
+        self.count_per_class: dict = defaultdict(int)
 
-        # Per-frame logs
-        self.detection_log: list[dict] = []
-        self.frame_stats:   list[dict] = []
+        # Logs
+        self.detection_log: list = []
+        self.frame_stats:   list = []
+
+        # Video writer (set in setup_video)
+        self.video_writer        = None
+        self.output_video_path   = None
 
         # COCO class-id filter
         self._class_ids = [
@@ -81,13 +90,23 @@ class TrafficTracker:
         ]
 
     # ── Setup ──────────────────────────────────────────────────────────────────
-    def setup_video(self, cap: cv2.VideoCapture):
+    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:
+            ts       = datetime.now().strftime("%Y%m%d_%H%M%S")
+            out_name = f"{self.scene_name}_{ts}_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[np.ndarray, dict]:
+    def process_frame(self, frame: np.ndarray) -> tuple:
         self.frame_index += 1
         timestamp = self.frame_index / self.fps
 
@@ -122,21 +141,16 @@ class TrafficTracker:
                 any_object_visible = True
                 per_class_in_frame[cls_name] += 1
 
-                # ── First-seen counting ───────────────────────────────────────
-                # Count each tracked object exactly once, the first frame it
-                # appears.  No line crossing required — every object in scene
-                # is counted as soon as it gets a stable track ID.
+                # First-seen counting
                 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
 
-                # Update trail history (tracked objects only)
                 if track_id != -1:
                     self.track_history[track_id].append((cx, cy))
 
-                # ── Log ───────────────────────────────────────────────────────
                 frame_detections.append({
                     "frame":      self.frame_index,
                     "timestamp":  round(timestamp, 3),
@@ -150,9 +164,8 @@ class TrafficTracker:
                 })
                 self.detection_log.append(frame_detections[-1])
 
-                # ── Draw bounding box ─────────────────────────────────────────
-                color  = CLASS_COLORS.get(cls_name, (255, 255, 255))
-                # Thicker border for newly counted objects (flash effect)
+                # Draw 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)
 
@@ -171,10 +184,12 @@ class TrafficTracker:
                         tc    = tuple(int(c * alpha) for c in color)
                         cv2.line(annotated, trail[i - 1], trail[i], tc, 2)
 
-        # ── Overlay counters (no counting line drawn) ─────────────────────────
         self._draw_counters(annotated, any_object_visible)
 
-        # ── Frame stats ───────────────────────────────────────────────────────
+        # Write annotated frame to output video
+        if self.video_writer is not None:
+            self.video_writer.write(annotated)
+
         frame_stat = {
             "frame":       self.frame_index,
             "timestamp":   round(timestamp, 3),
@@ -188,7 +203,7 @@ class TrafficTracker:
 
         return annotated, frame_stat
 
-    # ── Overlay ────────────────────────────────────────────────────────────────
+    # ── Counters overlay ───────────────────────────────────────────────────────
     def _draw_counters(self, frame: np.ndarray, any_visible: bool):
         overlay = frame.copy()
         cv2.rectangle(overlay, (0, 0),
@@ -198,13 +213,11 @@ class TrafficTracker:
 
         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))
-            label = f" {cls_name:<12} {count:>4}"
-            cv2.putText(frame, label, (8, y),
+            cv2.putText(frame, f" {cls_name:<12} {count:>4}", (8, y),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
             y += 22
 
@@ -212,25 +225,26 @@ class TrafficTracker:
             cv2.putText(frame, "NO OBJECTS DETECTED", (10, y + 10),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.5, (80, 80, 80), 1)
 
-    # ── Save logs ───────────────────────────────────────────────────────────────
+    # ── Finalise and save all logs ─────────────────────────────────────────────
     def save_logs(self) -> dict:
+        # Release video writer so the .mp4 is properly closed/flushed
+        if self.video_writer is not None:
+            self.video_writer.release()
+            self.video_writer = None
+
         ts     = datetime.now().strftime("%Y%m%d_%H%M%S")
         prefix = f"{self.scene_name}_{ts}"
 
-        # Detection log (JSONL)
         log_path = self.output_dir / f"{prefix}_detections.jsonl"
         with open(log_path, "w") as f:
             for entry in self.detection_log:
                 f.write(json.dumps(entry) + "\n")
 
-        # Summary JSON
         summary  = self.get_summary()
         sum_path = self.output_dir / f"{prefix}_summary.json"
         with open(sum_path, "w") as f:
             json.dump(summary, f, indent=2)
 
-        # Frame stats CSV
-        import csv
         csv_path = self.output_dir / f"{prefix}_frame_stats.csv"
         if self.frame_stats:
             with open(csv_path, "w", newline="") as f:
@@ -239,9 +253,10 @@ class TrafficTracker:
                 writer.writerows(self.frame_stats)
 
         return {
-            "detection_log": str(log_path),
-            "summary":       str(sum_path),
-            "frame_stats":   str(csv_path),
+            "detection_log":      str(log_path),
+            "summary":            str(sum_path),
+            "frame_stats":        str(csv_path),
+            "annotated_video":    self.output_video_path or "",
         }
 
     # ── Summary ─────────────────────────────────────────────────────────────────
@@ -249,23 +264,23 @@ class TrafficTracker:
         total_duration = self.frame_index / self.fps
         total_objects  = sum(self.count_per_class.values())
 
-        # Traffic intensity over time (10-second buckets)
-        buckets: dict[int, int] = defaultdict(int)
+        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,
-            "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),
+            "scene":                 self.scene_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,
         }

frontend/index.html

@@ -351,6 +351,7 @@
       <button class="btn btn-secondary" id="stopBtn" style="display:none" onclick="stopProcessing()">■ STOP</button>
       <div class="prog-bar"><div class="prog-fill" id="progFill"></div></div>
       <div class="prog-label" id="progLabel"></div>
+      <a class="btn btn-secondary" id="downloadVideoBtn" style="display:none;text-align:center;text-decoration:none;margin-top:7px" download>⬇ DOWNLOAD ANNOTATED VIDEO</a>
     </div>
 
     <!-- Total unique -->
@@ -639,6 +640,10 @@ function startSSE(sid) {
       setConn('live', 'DONE');
       document.getElementById('startBtn').disabled = false;
       document.getElementById('stopBtn').style.display = 'none';
+      // Show annotated video download button
+      const dlBtn = document.getElementById('downloadVideoBtn');
+      dlBtn.href  = `${API_BASE}/download/video/${currentSid}`;
+      dlBtn.style.display = 'block';
       return;
     }