Angle-based check

app.py

@@ -1,103 +1,111 @@
-import gradio as gr
-import numpy as np, cv2, json, tempfile, os
-from sklearn.cluster import KMeans
-
 # ============================================================
-# 🧩 1. Compute motion vectors from trajectory JSON
+# 🚦 Stage 3 – Wrong-Direction Detection (Angle + Temporal + Zones)
 # ============================================================
-def extract_motion_vectors(data):
-    vectors = []
-    for k, pts in data.items():
-        pts = np.array(pts)
-        if len(pts) < 2:
-            continue
-        diffs = np.diff(pts, axis=0)
-        for d in diffs:
-            if np.linalg.norm(d) > 1:   # ignore jitter / static points
-                vectors.append(d)
-    return np.array(vectors)
 
-
-# ============================================================
-# 🧮 2. Improved Dominant Flow Clustering (Cosine-based)
-# ============================================================
-def learn_flows_improved(vectors, n_clusters=2, normalize=True):
-    """
-    Improved dominant-flow clustering:
-    - Normalizes all vectors to unit direction (ignores speed)
-    - Clusters by angular orientation (cosine distance)
-    - Ignores low-magnitude / noisy motions
+import gradio as gr
+import numpy as np, cv2, json, os, tempfile
+from collections import defaultdict
+
+# ------------------------------------------------------------
+# ⚙️ CONFIG
+# ------------------------------------------------------------
+ANGLE_THRESHOLD = 60       # degrees → above this = wrong direction
+SMOOTH_FRAMES = 5          # frames used for temporal smoothing
+ENTRY_ZONE_RATIO = 0.15    # top 15% region = entry gate (skip labeling)
+
+# ------------------------------------------------------------
+# 1️⃣ Load flow model (from Stage 2 JSON or dict)
+# ------------------------------------------------------------
+def load_flow_model(flow_model_json):
+    """Expected keys:
+       {
+         "zones": N,
+         "zone_flow_centers": [[[dx,dy], ...], ...]
+       }
     """
-    if len(vectors) < n_clusters:
-        return None, None
-
-    # (1) Normalize to direction only
-    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
-    dirs = vectors / (norms + 1e-6)
+    model = json.load(open(flow_model_json))
+    centers = [np.array(z) for z in model["zone_flow_centers"]]
+    return centers
+
+# ------------------------------------------------------------
+# 2️⃣ Extract trajectories
+# ------------------------------------------------------------
+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
+
+# ------------------------------------------------------------
+# 3️⃣ Utility: compute smoothed direction
+# ------------------------------------------------------------
+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)
+
+# ------------------------------------------------------------
+# 4️⃣ Compute angle between two unit vectors
+# ------------------------------------------------------------
+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))
+
+# ------------------------------------------------------------
+# 5️⃣ Determine zone index for a y-coordinate
+# ------------------------------------------------------------
+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))
+
+# ------------------------------------------------------------
+# 6️⃣ 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)
 
-    # (2) Filter out tiny motions
-    valid = (norms[:, 0] > 1.5)
-    dirs = dirs[valid]
-    if len(dirs) < n_clusters:
-        return None, None
-
-    # (3) KMeans on direction vectors (≈ cosine distance)
-    kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=42)
-    kmeans.fit(dirs)
-    centers = kmeans.cluster_centers_
-
-    # (4) Normalize cluster centers again
-    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
-
-    # (5) Re-assign all original vectors to nearest angular center
-    sims = np.dot(vectors / (np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-6), centers.T)
-    labels = np.argmax(sims, axis=1)
-
-    return labels, centers
-
-
-# ============================================================
-# 🎨 3. Visualization Utility (Option A — Scaled-up Arrows)
-# ============================================================
-def draw_flow_overlay(vectors, labels, centers, bg_img=None):
-    # background
     if bg_img and os.path.exists(bg_img):
         bg = cv2.imread(bg_img)
-        if bg is None:
-            bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
     else:
         bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
+    h, w = bg.shape[:2]
 
     overlay = bg.copy()
-    colors = [(0, 0, 255), (255, 255, 0)]   # red & yellow
+    font = cv2.FONT_HERSHEY_SIMPLEX
 
-    # normalize arrow lengths for small samples
-    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
-    vectors = np.divide(vectors, norms + 1e-6) * 10
+    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))
 
-    # draw mini-arrows for field visualization
-    for i, ((vx, vy), lab) in enumerate(zip(vectors, labels)):
-        if i % 15 != 0:
+        # skip entry region
+        if y < h * ENTRY_ZONE_RATIO:
             continue
-        start = (np.random.randint(0, overlay.shape[1]),
-                 np.random.randint(0, overlay.shape[0]))
-        end = (int(start[0] + vx), int(start[1] + vy))
-        cv2.arrowedLine(overlay, start, end, colors[lab % 2], 1, tipLength=0.3)
-
-    # --- main dominant arrows ---
-    h, w = overlay.shape[:2]
-    scale = 300
-    center_pt = (w // 2, h // 2)
-
-    for i, c in enumerate(centers):
-        c = c / (np.linalg.norm(c) + 1e-6)
-        end = (int(center_pt[0] + c[0] * scale),
-               int(center_pt[1] + c[1] * scale))
-        offset = (i - 0.5) * 40
-        start = (center_pt[0], int(center_pt[1] + offset))
-        cv2.arrowedLine(overlay, start, end, (0, 255, 0), 4, tipLength=0.4)
-        cv2.putText(overlay, f"Flow {i+1}", (end[0] + 10, end[1]),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
+
+        # current smoothed direction
+        v = smooth_direction(pts)
+
+        # compare with zone's dominant flows (take best match)
+        centers = centers_by_zone[zone_idx]
+        angles = [angle_between(v, c) for c in centers]
+        best_angle = min(angles)
+        label = "OK" if best_angle < ANGLE_THRESHOLD else "WRONG"
+        color = (0, 255, 0) if label == "OK" else (0, 0, 255)
+
+        # draw trajectory and 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}", 
+                    (int(cur_pt[0]) + 5, int(cur_pt[1]) - 5),
+                    font, 0.6, color, 2)
 
     combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
     out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
@@ -105,59 +113,31 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None):
     return out_path
 
 
-# ============================================================
-# 🚀 4. Combined Pipeline
-# ============================================================
-def process_json(json_file, background=None):
-    try:
-        data = json.load(open(json_file))
-    except Exception as e:
-        return None, {"error": f"Invalid JSON file: {e}"}
-
-    vectors = extract_motion_vectors(data)
-    if len(vectors) == 0:
-        return None, {"error": "No motion vectors found."}
-
-    labels, centers = learn_flows_improved(vectors)
-    if labels is None:
-        return None, {"error": "Insufficient data for clustering."}
-
-    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
-    img_path = draw_flow_overlay(vectors, labels, centers, background)
-
-    stats = {
-        "num_vectors": int(len(vectors)),
-        "dominant_flows": int(len(centers)),
-        "flow_centers": centers.tolist()
-    }
-    return img_path, stats
-
-
-# ============================================================
-# 🖥️ 5. Gradio Interface
-# ============================================================
+# ------------------------------------------------------------
+# 🖥️ Gradio Interface
+# ------------------------------------------------------------
 description_text = """
-### 🧭 Dominant Flow Learning (Stage 2 — Cosine-Based Improved)
-Upload the **trajectories JSON** from Stage 1.  
-Optionally upload a background frame for overlay visualization.
+### 🚦 Wrong-Direction Detection (Stage 3 — Angle + Temporal + Zone-Aware)
+1. Upload the **trajectories JSON** from Stage 1.  
+2. Upload the **flow model JSON** from Stage 2 (zone-wise).  
+3. Optionally add a background road frame for overlay.  
+
+**Logic:**  
+- Compares each vehicle’s smoothed direction vector with the dominant flow of its zone.  
+- Ignores top-entry region to avoid false positives.  
+- Flags vehicles as WRONG if angular difference > 60°.
 """
 
-example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
-example_bg = "frame_sample.jpg" if os.path.exists("frame_sample.jpg") else None
-
 demo = gr.Interface(
-    fn=process_json,
+    fn=classify_wrong_direction,
     inputs=[
-        gr.File(label="Upload trajectories JSON"),
+        gr.File(label="Trajectories JSON (Stage 1)"),
+        gr.File(label="Flow Model JSON (Stage 2)"),
         gr.File(label="Optional background frame (.jpg)")
     ],
-    outputs=[
-        gr.Image(label="Dominant Flow Overlay"),
-        gr.JSON(label="Flow Stats")
-    ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Cosine-Based Improved)",
-    description=description_text,
-    examples=[[example_json, example_bg]] if example_json else None,
+    outputs=gr.Image(label="Wrong-Direction Overlay"),
+    title="🚗 Wrong-Direction Detection – Stage 3 (Angle + Temporal + Zones)",
+    description=description_text
 )
 
 if __name__ == "__main__":