Stage 3 (Angle + Temporal + Zone-Aware + Entry Gating)

app.py

@@ -1,211 +1,167 @@
-import os, cv2, json, tempfile, zipfile, numpy as np, gradio as gr
-from ultralytics import YOLO
-from filterpy.kalman import KalmanFilter
-from scipy.optimize import linear_sum_assignment
+# ============================================================
+# 🚦 Stage 3 – Wrong-Direction Detection
+#     (Angle + Temporal + Zone-Aware + Entry Gating + Confidence)
+# ============================================================
+
+import gradio as gr
+import numpy as np, cv2, json, os, tempfile
+from collections import defaultdict
 
 # ------------------------------------------------------------
-# 🔧  Safe-load fix for PyTorch 2.6
+# ⚙️ CONFIG
 # ------------------------------------------------------------
-import torch, ultralytics.nn.tasks as ultralytics_tasks
-torch.serialization.add_safe_globals([ultralytics_tasks.DetectionModel])
+ANGLE_THRESHOLD = 60        # degrees → above this = WRONG
+SMOOTH_FRAMES = 5           # frames for temporal smoothing
+ENTRY_ZONE_RATIO = 0.15     # top 15% = entry region (skip)
+CONF_MIN, CONF_MAX = 0, 100
 
 # ------------------------------------------------------------
-# ⚙️  YOLO setup
+# 1️⃣ Load flow model (Stage 2)
 # ------------------------------------------------------------
-MODEL_PATH = "yolov8n.pt"
-model = YOLO(MODEL_PATH)
-VEHICLE_CLASSES = [2, 3, 5, 7]  # car, motorcycle, bus, truck
-
+def load_flow_model(flow_model_json):
+    model = json.load(open(flow_model_json))
+    centers = [np.array(z) for z in model["zone_flow_centers"]]
+    return centers
 
 # ------------------------------------------------------------
-# 🧩  Kalman tracker
+# 2️⃣ Extract trajectories
 # ------------------------------------------------------------
-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] = np.array(self.centroid(bbox)).reshape(2,1)
-        self.trace = []
-
-    def centroid(self, b):
-        x1, y1, x2, y2 = b
-        return [(x1+x2)/2, (y1+y2)/2]
-
-    def predict(self):
-        self.kf.predict()
-        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)))
-        return (cx, cy)
-
+def extract_trajectories(json_file):
+    data = json.load(open(json_file))
+    tracks = {tid: np.array(pts) for tid, pts in data.items() if len(pts) > 2}
+    return tracks
 
 # ------------------------------------------------------------
-# 🧮  Direction analyzer
+# 3️⃣ Smoothed direction for a trajectory
 # ------------------------------------------------------------
-def analyze_direction(trace, centers):
-    if len(trace) < 3:
-        return "NA", 1.0
-    v = np.array(trace[-1]) - np.array(trace[-3])
-    if np.linalg.norm(v) < 1e-6:
-        return "NA", 1.0
-    v = v / np.linalg.norm(v)
-    sims = np.dot(centers, v)
-    max_sim = np.max(sims)
-    if max_sim < 0:
-        return "WRONG", float(max_sim)
-    return "OK", float(max_sim)
-
+def smooth_direction(pts, window=SMOOTH_FRAMES):
+    if len(pts) < 2:
+        return np.array([0, 0])
+    diffs = np.diff(pts[-window:], axis=0)
+    v = np.mean(diffs, axis=0)
+    n = np.linalg.norm(v)
+    return v / (n + 1e-6)
 
 # ------------------------------------------------------------
-# 🧭  Load normalized flow centers
+# 4️⃣ Compute angular difference (deg)
 # ------------------------------------------------------------
-def load_flow_centers(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)
-    return centers
-
+def angle_between(v1, v2):
+    v1 = v1 / (np.linalg.norm(v1) + 1e-6)
+    v2 = v2 / (np.linalg.norm(v2) + 1e-6)
+    cosang = np.clip(np.dot(v1, v2), -1, 1)
+    return np.degrees(np.arccos(cosang))
 
 # ------------------------------------------------------------
-# 🎥  Process video
-# ------------------------------------------------------------
-def process_video(video_path, flow_json, show_only_wrong=False):
-    centers = load_flow_centers(flow_json)
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS) or 25
-    w, h = int(cap.get(3)), int(cap.get(4))
-
-    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
-
-    tracks, next_id, log = [], 0, []
-
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        results = model(frame, verbose=False)[0]
-        detections = []
-        for box in results.boxes:
-            if int(box.cls) in VEHICLE_CLASSES and box.conf > 0.3:
-                detections.append(box.xyxy[0].cpu().numpy())
-
-        # Predict existing
-        predicted = [t.predict() for t in tracks]
-        predicted = np.array(predicted) if len(predicted) > 0 else np.empty((0,2))
-
-        # Assign detections to tracks
-        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)
-                next_id += 1
-                t.update(d)
-                tracks.append(t)
-
-        # --- 🧩 Draw + Log (toggle support) ---
-        for trk in tracks:
-            if len(trk.trace) < 3:
-                continue
-            status, sim = analyze_direction(trk.trace, centers)
-
-            # Skip OKs if toggle is enabled
-            if show_only_wrong and status != "WRONG":
-                continue
-
-            x, y = map(int, trk.trace[-1])
-            color = (0,255,0) if status=="OK" else ((0,0,255) if status=="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)
-
-            # Log once per unique vehicle
-            if len(trk.trace) > 5 and not any(entry["id"] == trk.id for entry in log):
-                log.append({"id": trk.id, "status": status, "cos_sim": round(sim,3)})
-
-        out.write(frame)
-
-    cap.release()
-    out.release()
-
-    # Unique summary
-    unique_ids = {entry["id"] for entry in log}
-    summary = {"vehicles_analyzed": len(unique_ids)}
-
-    # Create ZIP bundle
-    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
-
+# 5️⃣ Determine zone index for y
+# ------------------------------------------------------------
+def get_zone_idx(y, frame_h, n_zones):
+    zone_height = frame_h / n_zones
+    return int(np.clip(y // zone_height, 0, n_zones - 1))
 
 # ------------------------------------------------------------
-# 🖥️  Gradio interface
+# 6️⃣ Confidence mapping
 # ------------------------------------------------------------
-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
+def angle_to_confidence(angle):
+    """
+    0° → 100% confidence
+    ANGLE_THRESHOLD° → 50%
+    180° → 0%
+    """
+    if angle < 0:
+        return CONF_MIN
+    if angle >= 180:
+        return CONF_MIN
+    # linear mapping: smaller angle = higher confidence
+    conf = max(CONF_MIN, CONF_MAX - (angle / 180) * 100)
+    return round(conf, 1)
 
+# ------------------------------------------------------------
+# 7️⃣ Main logic
+# ------------------------------------------------------------
+def classify_wrong_direction(traj_json, flow_model_json, bg_img=None):
+    tracks = extract_trajectories(traj_json)
+    centers_by_zone = load_flow_model(flow_model_json)
+
+    if bg_img and os.path.exists(bg_img):
+        bg = cv2.imread(bg_img)
+    else:
+        bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
+    h, w = bg.shape[:2]
+
+    overlay = bg.copy()
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    results = []
+
+    for tid, pts in tracks.items():
+        if len(pts) < 3:
+            continue
+        cur_pt = pts[-1]
+        y = cur_pt[1]
+        zone_idx = get_zone_idx(y, h, len(centers_by_zone))
+
+        # Skip entry region
+        if y < h * ENTRY_ZONE_RATIO:
+            continue
+
+        v = smooth_direction(pts)
+        centers = centers_by_zone[zone_idx]
+        angles = [angle_between(v, c) for c in centers]
+        best_angle = min(angles)
+
+        # Confidence & label
+        conf = angle_to_confidence(best_angle)
+        label = "OK" if best_angle < ANGLE_THRESHOLD else "WRONG"
+        color = (0, 255, 0) if label == "OK" else (0, 0, 255)
+
+        # Draw trajectory & label
+        for p1, p2 in zip(pts[:-1], pts[1:]):
+            cv2.line(overlay, tuple(p1.astype(int)), tuple(p2.astype(int)), color, 2)
+        cv2.circle(overlay, tuple(cur_pt.astype(int)), 5, color, -1)
+        cv2.putText(
+            overlay,
+            f"ID:{tid} {label} ({conf}%)",
+            (int(cur_pt[0]) + 5, int(cur_pt[1]) - 5),
+            font, 0.6, color, 2
+        )
+
+        results.append({
+            "id": tid,
+            "zone": int(zone_idx),
+            "angle": round(best_angle, 1),
+            "confidence": conf,
+            "label": label
+        })
+
+    combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
+    out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
+    cv2.imwrite(out_path, combined)
+    return out_path, results
 
+# ------------------------------------------------------------
+# 🖥️ Gradio Interface
+# ------------------------------------------------------------
 description_text = """
-### 🚦 Wrong-Direction Detection (Stage 3)
-Upload your traffic video and the **flow_stats.json** from Stage 2.  
-You can toggle whether to display all detections or only WRONG-direction vehicles.
+### 🚦 Wrong-Direction Detection (Stage 3 — with Confidence)
+- Compares each vehicle’s motion to its zone’s dominant flow.  
+- Uses angular difference → smaller angle ⇒ higher confidence.  
+- Ignores entry region to avoid false positives.  
+- Displays ID, label, and confidence percentage.
 """
 
 demo = gr.Interface(
-    fn=run_app,
+    fn=classify_wrong_direction,
     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)
+        gr.File(label="Trajectories JSON (Stage 1)"),
+        gr.File(label="Flow Model JSON (Stage 2)"),
+        gr.File(label="Optional background frame (.jpg)")
     ],
     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.Image(label="Annotated Output"),
+        gr.JSON(label="Per-Vehicle Results")
     ],
-    title="🚗 Wrong-Direction Detection – Stage 3 (Toggle + ZIP)",
-    description=description_text,
-    examples=None,
+    title="🚗 Stage 3 — Wrong-Direction Detection (with Confidence)",
+    description=description_text
 )
 
-# Disable analytics / flagging / SSR
-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()