speed + confidence labels on the output video

app.py

@@ -1,96 +1,184 @@
 # ============================================================
-# 🚗 Stage 4 — Speed Calculation (Gradio UI)
+# 🚗 Stage 4 — Speed Calculation (Video + Confidence + Filter)
 # ============================================================
 
 import gradio as gr
-import numpy as np, json, cv2, tempfile, os, math, time
-from collections import defaultdict, deque
+import cv2, os, json, tempfile, math, time
+import numpy as np
+from ultralytics import YOLO
+from filterpy.kalman import KalmanFilter
+from scipy.optimize import linear_sum_assignment
 
 # ------------------------------------------------------------
-# 🧠 Configuration defaults
+# 🧠 Safe-load fix for PyTorch 2.6
 # ------------------------------------------------------------
-DEFAULT_PIXEL_TO_METER = 0.05     # meters per pixel
-DEFAULT_SPEED_LIMIT = 60.0        # km/h
-WINDOW_SIZE = 5                   # frames for moving average
+import torch, ultralytics.nn.tasks as ultralytics_tasks
+torch.serialization.add_safe_globals([ulralytics_tasks.DetectionModel])
 
 # ------------------------------------------------------------
-# 🧮 Compute speeds from trajectories JSON
+# ⚙️ Model + Config
 # ------------------------------------------------------------
-def compute_speeds(json_file, pixel_to_meter, speed_limit):
-    try:
-        data = json.load(open(json_file))
-    except Exception as e:
-        return None, {"error": f"Invalid JSON file: {e}"}
-
-    # Input format: {track_id: [[x,y], [x,y], ...]}
-    results = {}
-    overlay = np.ones((600, 900, 3), dtype=np.uint8) * 40
-
-    for tid, pts in data.items():
-        pts = np.array(pts, dtype=float)
-        if len(pts) < 2:
-            continue
-
-        # compute displacement per frame
-        diffs = np.diff(pts, axis=0)
-        dists_pix = np.linalg.norm(diffs, axis=1)
-        # assume 30 FPS default → Δt = 1/30 s
-        speeds_m_s = (dists_pix * pixel_to_meter) * 30.0
-        speeds_kmph = speeds_m_s * 3.6
-        avg_speed = np.mean(speeds_kmph)
-        status = "SPEEDING" if avg_speed > speed_limit else "OK"
-
-        results[str(tid)] = {
-            "avg_speed_kmph": round(float(avg_speed), 2),
-            "max_speed_kmph": round(float(np.max(speeds_kmph)), 2),
-            "status": status
-        }
-
-
-        # draw on overlay (simple visualization)
-        color = (0, 0, 255) if status == "SPEEDING" else (0, 255, 0)
-        start = tuple(np.int32(pts[0]))
-        end = tuple(np.int32(pts[-1]))
-        cv2.arrowedLine(overlay, start, end, color, 2, tipLength=0.2)
-        cv2.putText(
-            overlay,
-            f"ID:{tid} {avg_speed:.1f}km/h",
-            end,
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.6,
-            color,
-            2,
-        )
-
-    # Save overlay
-    out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
-    cv2.imwrite(out_path, overlay)
-
-    return out_path, results
+MODEL_PATH = "yolov8n.pt"
+model = YOLO(MODEL_PATH)
+VEHICLE_CLASSES = [2, 3, 5, 7]  # car, motorcycle, bus, truck
 
+# Speed calculation constants
+PIXEL_TO_METER = 0.05
+FPS_DEFAULT = 30.0
 
-# ------------------------------------------------------------
-# 🖥️  Gradio Interface
-# ------------------------------------------------------------
+# ============================================================
+# 🧩 Kalman-based Tracker
+# ============================================================
+class Track:
+    def __init__(self, bbox, tid):
+        self.id = tid
+        self.kf = KalmanFilter(dim_x=4, dim_z=2)
+        self.kf.F = np.array([[1,0,1,0],[0,1,0,1],[0,0,1,0],[0,0,0,1]])
+        self.kf.H = np.array([[1,0,0,0],[0,1,0,0]])
+        self.kf.P *= 10
+        self.kf.R *= 1
+        self.kf.x[:2] = np.array(bbox[:2]).reshape(2,1)
+        self.history = []
+        self.frames_seen = 0
+        self.avg_speed = 0
+        self.confidence = 1.0
+        self.status = "OK"
+
+    def update(self, bbox):
+        self.kf.predict()
+        self.kf.update(np.array(bbox[:2]))
+        x, y = self.kf.x[:2].reshape(-1)
+        self.history.append([x, y])
+        if len(self.history) > 30:
+            self.history.pop(0)
+        self.frames_seen += 1
+        return [x, y]
+
+
+# ============================================================
+# 🧮 Utility Functions
+# ============================================================
+def compute_speed(track, fps, pixel_to_meter):
+    if len(track.history) < 2:
+        return 0.0
+    pts = np.array(track.history)
+    diffs = np.diff(pts, axis=0)
+    dists = np.linalg.norm(diffs, axis=1)
+    mean_pix_per_frame = np.mean(dists)
+    speed_m_s = mean_pix_per_frame * pixel_to_meter * fps
+    return speed_m_s * 3.6  # km/h
+
+
+# ============================================================
+# 🚀 Main Processing Function
+# ============================================================
+def process_video(video_file, speed_limit, pixel_to_meter, confidence_filter, show_only_speeding):
+
+    cap = cv2.VideoCapture(video_file)
+    fps = cap.get(cv2.CAP_PROP_FPS) or FPS_DEFAULT
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+
+    tracks, next_id = {}, 0
+    frame_no = 0
+    DELAY_FRAMES = 5
+    SPEED_SMOOTH_ALPHA = 0.5  # exponential moving average
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_no += 1
+        results = model(frame)[0]
+        dets = []
+        for box in results.boxes:
+            cls = int(box.cls[0])
+            conf = float(box.conf[0])
+            if cls in VEHICLE_CLASSES:
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+                cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+                dets.append([cx, cy, conf])
+        dets = np.array(dets)
+
+        # --- Tracker update ---
+        assigned = set()
+        if len(dets) > 0 and len(tracks) > 0:
+            existing = np.array([t.kf.x[:2].reshape(-1) for t in tracks.values()])
+            dists = np.linalg.norm(existing[:, None, :] - dets[None, :, :2], axis=2)
+            row_idx, col_idx = linear_sum_assignment(dists)
+            for r, c in zip(row_idx, col_idx):
+                if dists[r, c] < 60:
+                    tid = list(tracks.keys())[r]
+                    tracks[tid].update(dets[c])
+                    tracks[tid].confidence = float(dets[c][2])
+                    assigned.add(c)
+        for i, d in enumerate(dets):
+            if i not in assigned:
+                tracks[next_id] = Track(d, next_id)
+                tracks[next_id].confidence = float(d[2])
+                next_id += 1
+
+        # --- Speed & Draw ---
+        for tid, trk in list(tracks.items()):
+            pos = trk.update(trk.kf.x[:2].reshape(-1))
+            if trk.frames_seen < DELAY_FRAMES:
+                continue
+
+            # compute speed
+            speed = compute_speed(trk, fps, pixel_to_meter)
+            # smooth speed
+            trk.avg_speed = SPEED_SMOOTH_ALPHA * speed + (1 - SPEED_SMOOTH_ALPHA) * trk.avg_speed
+            status = "SPEEDING" if trk.avg_speed > speed_limit else "OK"
+            trk.status = status
+
+            # skip by confidence filter
+            if trk.confidence < confidence_filter:
+                continue
+            if show_only_speeding and trk.status != "SPEEDING":
+                continue
+
+            color = (0, 0, 255) if status == "SPEEDING" else (0, 255, 0)
+            label = f"ID:{tid} {trk.avg_speed:.1f}km/h ({trk.confidence:.2f})"
+            cv2.putText(frame, label, tuple(np.int32(pos)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+            cv2.circle(frame, tuple(np.int32(pos)), 4, color, -1)
+
+        out.write(frame)
+
+    cap.release()
+    out.release()
+    return out_path
+
+
+# ============================================================
+# 🎛️ Gradio UI
+# ============================================================
 description = """
-### 🧮 Stage 4 — Speed Calculation
-Uploads the **trajectories JSON** (same format as Stage 1 output).  
-Calculates per-vehicle speed (km/h) using a pixel-to-meter calibration and 30 FPS default.  
-Color-codes SPEEDING vehicles (red).  
+### 🚗 Stage 4 — Speed Calculation (Video + Confidence + Filtering)
+Uploads a traffic video, detects and tracks vehicles,  
+computes their **approximate speed**, and overlays **speed + confidence labels**.
+
+**Controls:**
+- 🎚️ Pixel → Meter conversion for calibration  
+- 🚧 Speed limit for violation tagging  
+- 🧠 Confidence threshold (hide low-confidence detections)  
+- 🚨 Option to show only SPEEDING vehicles
 """
 
 demo = gr.Interface(
-    fn=compute_speeds,
+    fn=process_video,
     inputs=[
-        gr.File(label="Upload trajectories JSON"),
-        gr.Slider(0.01, 0.2, value=DEFAULT_PIXEL_TO_METER, step=0.01, label="Pixel → Meter Conversion"),
-        gr.Slider(10, 120, value=DEFAULT_SPEED_LIMIT, step=5, label="Speed Limit (km/h)"),
+        gr.File(label="Upload Traffic Video (.mp4)"),
+        gr.Slider(10, 120, value=60, step=5, label="Speed Limit (km/h)"),
+        gr.Slider(0.01, 0.2, value=0.05, step=0.01, label="Pixel → Meter Conversion"),
+        gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Confidence Filter (Show ≥ this value)"),
+        gr.Checkbox(label="Show ONLY Speeding Vehicles", value=False)
     ],
     outputs=[
-        gr.Image(label="Speed Overlay"),
-        gr.JSON(label="Speed Stats (Stage 4 Output)")
+        gr.Video(label="Output Video (Speed Overlay)")
     ],
-    title="🚗 Stage 4 – Speed Calculation",
+    title="🚗 Stage 4 – Speed Calculation with Confidence & Filter",
     description=description,
 )