Road curvature

app.py

@@ -1,269 +1,173 @@
-import os, cv2, json, tempfile, zipfile, numpy as np, gradio as gr
+# ============================================================
+# 🚦 Stage 3 — Wrong Direction Detection (Improved)
+# ============================================================
+
+import os, cv2, json, tempfile, numpy as np, gradio as gr
 from ultralytics import YOLO
 from filterpy.kalman import KalmanFilter
 from scipy.optimize import linear_sum_assignment
 
 # ------------------------------------------------------------
-# 🔧  Safe-load fix for PyTorch 2.6
+# 🧠 Safe-load fix for PyTorch 2.6
 # ------------------------------------------------------------
 import torch, ultralytics.nn.tasks as ultralytics_tasks
 torch.serialization.add_safe_globals([ultralytics_tasks.DetectionModel])
 
-# ------------------------------------------------------------
-# ⚙️  YOLO setup
-# ------------------------------------------------------------
 MODEL_PATH = "yolov8n.pt"
 model = YOLO(MODEL_PATH)
 VEHICLE_CLASSES = [2, 3, 5, 7]  # car, motorcycle, bus, truck
 
-# ------------------------------------------------------------
-# 🧭  Utility: rotate vector by angle
-# ------------------------------------------------------------
-def rotate_vec(v, theta_deg):
-    t = np.deg2rad(theta_deg)
-    R = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]])
-    return R @ v
-
-# ------------------------------------------------------------
-# 🧩  Kalman tracker with temporal smoothing + entry gate flag
-# ------------------------------------------------------------
+# ============================================================
+# 🧩 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 *= 1000.0
-        self.kf.R *= 10.0
-        self.kf.x[:2,0] = np.array(self.centroid(bbox), dtype=float)
-        self.trace = []
-        self.status_hist = []
-        self.entry_flag = False
-        self.active = True
-        self.missed_frames = 0
-
-    def centroid(self, b):
-        x1, y1, x2, y2 = b
-        return [(x1+x2)/2, (y1+y2)/2]
-
-    def predict(self):
+        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.status = "OK"
+
+    def update(self, bbox):
         self.kf.predict()
-        self.missed_frames += 1
-        return self.kf.x[:2].reshape(2)
-
-    def update(self, b):
-        z = np.array(self.centroid(b)).reshape(2,1)
-        self.kf.update(z)
-        cx, cy = self.kf.x[:2].reshape(2)
-        self.trace.append((float(cx), float(cy)))
-        self.missed_frames = 0
-        return (cx, cy)
-
-# ------------------------------------------------------------
-# 🧮  Direction analyzer (angle + temporal aware)
-# ------------------------------------------------------------
-def analyze_direction(trace, centers, road_angle_deg, hist):
-    if len(trace) < 3:
-        return "NA", 1.0
-
-    # motion vector
-    v = np.array(trace[-1]) - np.array(trace[-3])
-    if np.linalg.norm(v) < 1e-6:
-        return "NA", 1.0
-
-    # rotate to road reference
-    v = rotate_vec(v / np.linalg.norm(v), -road_angle_deg)
-
-    # cosine similarity vs dominant centers
-    sims = np.dot(centers, v)
-    max_sim = np.max(sims)
-
-    # temporal averaging
-    hist.append(max_sim)
-    if len(hist) > 5:
-        hist.pop(0)
-    avg_sim = np.mean(hist)
-
-    if avg_sim < -0.2:
-        return "WRONG", float(avg_sim)
-    elif avg_sim > 0.2:
-        return "OK", float(avg_sim)
-    else:
-        return "NA", float(avg_sim)
-
-# ------------------------------------------------------------
-# 🗺️  Load Stage-2 flow stats (centers, angle, zones)
-# ------------------------------------------------------------
-def load_flow_stats(flow_json):
-    data = json.load(open(flow_json))
-    centers = np.array(data["flow_centers"])
-    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
-    road_angle_deg = float(data.get("road_angle_deg", 0.0))
-    drive_zone = data.get("drive_zone", None)
-    entry_zones = data.get("entry_zones", [])
-    return centers, road_angle_deg, drive_zone, entry_zones
-
-# ------------------------------------------------------------
-# 🧾  Zone tests
-# ------------------------------------------------------------
-def inside_zone(pt, zone):
-    if zone is None: return True
-    return cv2.pointPolygonTest(np.array(zone, np.int32), pt, False) >= 0
-
-def inside_any(pt, zones):
-    return any(cv2.pointPolygonTest(np.array(z, np.int32), pt, False) >= 0 for z in zones)
-
-# ------------------------------------------------------------
-# 🎥  Process video (angle + temporal + zone + entry-gating)
-# ------------------------------------------------------------
-def process_video(video_path, flow_json, show_only_wrong=False):
-    centers, road_angle, drive_zone, entry_zones = load_flow_stats(flow_json)
-
-    cap = cv2.VideoCapture(video_path)
-    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 25
-    w, h = int(cap.get(3)), int(cap.get(4))
-
+        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]
+
+# ============================================================
+# ⚙️ Utilities
+# ============================================================
+def compute_cosine_similarity(v1, v2):
+    v1 = v1 / (np.linalg.norm(v1) + 1e-6)
+    v2 = v2 / (np.linalg.norm(v2) + 1e-6)
+    return np.dot(v1, v2)
+
+def smooth_direction(points, window=5):
+    """Compute smoothed motion vector using last N points"""
+    if len(points) < window + 1:
+        return None
+    diffs = np.diff(points[-window:], axis=0)
+    avg_vec = np.mean(diffs, axis=0)
+    if np.linalg.norm(avg_vec) < 1:
+        return None
+    return avg_vec
+
+# ============================================================
+# 🧭 Wrong-Direction Detection Core
+# ============================================================
+def process_video(video_file, stage2_json):
+    data = json.load(open(stage2_json))
+    lane_flows = np.array(data.get("flow_centers", [[1,0]]))
+    drive_zone = np.array(data.get("drive_zone", []))
+    entry_zones = [np.array(z) for z in data.get("entry_zones", [])]
+
+    cap = cv2.VideoCapture(video_file)
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), 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))
+    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
 
-    tracks, next_id, log = [], 0, []
+    tracks, next_id = {}, 0
+    SIM_THRESH = 0.5          # cosine similarity threshold
+    DELAY_FRAMES = 8          # wait N frames before flagging
+    MIN_FLOW_SPEED = 1.2      # ignore jitter
 
     while True:
         ret, frame = cap.read()
-        if not ret: break
+        if not ret:
+            break
 
-        results = model(frame, verbose=False)[0]
-        detections = []
+        results = model(frame)[0]
+        dets = []
         for box in results.boxes:
-            if int(box.cls) in VEHICLE_CLASSES and float(box.conf) > 0.3:
-                detections.append(box.xyxy[0].cpu().numpy())
-
-        # predict existing
-        predicted = [t.predict() for t in tracks if t.active]
-        predicted = np.array(predicted) if predicted else np.empty((0,2))
-
-        # assign detections
+            cls = int(box.cls[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])
+        dets = np.array(dets)
+
+        # --- Tracker update ---
         assigned = set()
-        if len(predicted) > 0 and len(detections) > 0:
-            cost = np.zeros((len(predicted), len(detections)))
-            for i, p in enumerate(predicted):
-                for j, d in enumerate(detections):
-                    cx, cy = ((d[0]+d[2])/2, (d[1]+d[3])/2)
-                    cost[i,j] = np.linalg.norm(p - np.array([cx,cy]))
-            r, c = linear_sum_assignment(cost)
-            for i, j in zip(r, c):
-                if cost[i,j] < 80:
-                    assigned.add(j)
-                    tracks[i].update(detections[j])
-
-        # new tracks
-        for j, d in enumerate(detections):
-            if j not in assigned:
-                t = Track(d, next_id)
+        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, :, :], axis=2)
+            row_idx, col_idx = linear_sum_assignment(dists)
+            for r, c in zip(row_idx, col_idx):
+                if dists[r, c] < 50:
+                    tid = list(tracks.keys())[r]
+                    tracks[tid].update(dets[c])
+                    assigned.add(c)
+        for i, d in enumerate(dets):
+            if i not in assigned:
+                tracks[next_id] = Track(d, next_id)
                 next_id += 1
-                t.update(d)
-                first_pt = tuple(map(int, t.trace[-1]))
-                # entry gating: mark if starts inside forbidden zone
-                if inside_any(first_pt, entry_zones):
-                    t.entry_flag = True
-                tracks.append(t)
 
-        # clean up stale tracks
-        for t in tracks:
-            if t.missed_frames > 15:
-                t.active = False
-
-        # draw + log
-        for trk in tracks:
-            if not trk.active or len(trk.trace) < 3:
+        # --- Draw & classify ---
+        for tid, trk in list(tracks.items()):
+            pos = trk.update(trk.kf.x[:2].reshape(-1))
+            pts = np.array(trk.history)
+            if len(pts) > 1:
+                for i in range(1, len(pts)):
+                    cv2.line(frame, tuple(np.int32(pts[i-1])), tuple(np.int32(pts[i])), (0, 0, 255), 1)
+
+            # compute smooth direction
+            motion = smooth_direction(pts)
+            if motion is None:
                 continue
-
-            x, y = map(int, trk.trace[-1])
-            if not inside_zone((x, y), drive_zone):
-                continue  # skip outside drive zone
-
-            status, sim = analyze_direction(trk.trace, centers, road_angle, trk.status_hist)
-            if trk.entry_flag:
-                status = "WRONG_ENTRY"
-
-            if show_only_wrong and status not in ["WRONG", "WRONG_ENTRY"]:
+            if np.linalg.norm(motion) < MIN_FLOW_SPEED:
                 continue
 
-            color = (0,255,0) if status=="OK" else \
-                    (0,0,255) if status.startswith("WRONG") else (200,200,200)
-
-            cv2.circle(frame, (x,y), 4, color, -1)
-            cv2.putText(frame, f"ID:{trk.id} {status}", (x-20,y-10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
-            for i in range(1, len(trk.trace)):
-                cv2.line(frame,
-                         (int(trk.trace[i-1][0]), int(trk.trace[i-1][1])),
-                         (int(trk.trace[i][0]), int(trk.trace[i][1])),
-                         color, 1)
-
-            if len(trk.trace) > 5 and not any(e["id"]==trk.id for e in log):
-                log.append({
-                    "id": trk.id,
-                    "status": status,
-                    "cos_sim": round(sim,3),
-                    "entry_flag": trk.entry_flag
-                })
+            # cosine similarity to closest lane flow
+            sims = [compute_cosine_similarity(motion, f) for f in lane_flows]
+            best_sim = max(sims)
+
+            # only classify after some frames (to reduce false early flag)
+            if trk.frames_seen > DELAY_FRAMES:
+                if best_sim < SIM_THRESH:
+                    trk.status = "WRONG"
+                    color = (0, 0, 255)
+                else:
+                    trk.status = "OK"
+                    color = (0, 255, 0)
+                cv2.putText(frame, f"ID:{tid} {trk.status}", tuple(np.int32(pos)),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
         out.write(frame)
 
     cap.release()
     out.release()
-
-    # summary
-    unique_ids = {e["id"] for e in log}
-    summary = {
-        "vehicles_analyzed": len(unique_ids),
-        "wrong_count": sum(1 for e in log if e["status"].startswith("WRONG")),
-        "road_angle_deg": road_angle
-    }
-
-    # zip outputs
-    zip_path = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
-    with zipfile.ZipFile(zip_path, "w") as zf:
-        zf.write(out_path, arcname="violation_output.mp4")
-        zf.writestr("per_vehicle_log.json", json.dumps(log, indent=2))
-        zf.writestr("summary.json", json.dumps(summary, indent=2))
-
-    return out_path, log, summary, zip_path
-
-# ------------------------------------------------------------
-# 🖥️  Gradio interface
-# ------------------------------------------------------------
-def run_app(video, flow_file, show_only_wrong):
-    vid, log_json, summary, zip_file = process_video(video, flow_file, show_only_wrong)
-    return vid, log_json, summary, zip_file
-
-description_text = """
-### 🚦 Wrong-Direction Detection (Stage 3 — Angle + Temporal + Zone + Entry-Aware)
-Upload your traffic video and the **flow_stats.json** from Stage 2.  
-Stage 3 will respect the learned road angle, driving zones, and entry gates.
+    return out_path
+
+# ============================================================
+# 🎛️ Gradio Interface
+# ============================================================
+description = """
+### 🚦 Stage 3 — Wrong Direction Detection (Improved)
+- Uses cosine similarity instead of raw angle comparison  
+- Lane-wise flow support for curved roads  
+- Temporal smoothing & delayed classification  
 """
 
 demo = gr.Interface(
-    fn=run_app,
+    fn=process_video,
     inputs=[
-        gr.Video(label="Upload Traffic Video (.mp4)"),
-        gr.File(label="Upload flow_stats.json (Stage 2 Output)"),
-        gr.Checkbox(label="Show Only Wrong Labels", value=False)
-    ],
-    outputs=[
-        gr.Video(label="Violation Output Video"),
-        gr.JSON(label="Per-Vehicle Log"),
-        gr.JSON(label="Summary"),
-        gr.File(label="⬇️ Download All Outputs (ZIP)")
+        gr.File(label="Input Video"),
+        gr.File(label="Stage 2 Flow JSON")
     ],
-    title="🚗 Wrong-Direction Detection – Stage 3 (Angle + Temporal + Zone + Entry)",
-    description=description_text,
+    outputs=gr.Video(label="Output (with WRONG/OK labels)"),
+    title="🚗 Stage 3 – Improved Wrong-Direction Detection",
+    description=description
 )
 
-demo.flagging_mode = "never"
-demo.cache_examples = False
-os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False, show_api=False)
+    demo.launch()