revert

app.py

@@ -1,111 +1,103 @@
+import gradio as gr
+import numpy as np, cv2, json, tempfile, os
+from sklearn.cluster import KMeans
+
 # ============================================================
-# 🚦 Stage 3 – Wrong-Direction Detection (Angle + Temporal + Zones)
+# 🧩 1. Compute motion vectors from trajectory JSON
 # ============================================================
+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)
 
-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], ...], ...]
-       }
+
+# ============================================================
+# 🧮 2. Improved Dominant Flow Clustering (Cosine-based)
+# ============================================================
+def learn_flows_improved(vectors, n_clusters=2, normalize=True):
     """
-    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)
+    Improved dominant-flow clustering:
+    - Normalizes all vectors to unit direction (ignores speed)
+    - Clusters by angular orientation (cosine distance)
+    - Ignores low-magnitude / noisy motions
+    """
+    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)
 
+    # (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()
-    font = cv2.FONT_HERSHEY_SIMPLEX
+    colors = [(0, 0, 255), (255, 255, 0)]   # red & yellow
 
-    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))
+    # normalize arrow lengths for small samples
+    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+    vectors = np.divide(vectors, norms + 1e-6) * 10
 
-        # skip entry region
-        if y < h * ENTRY_ZONE_RATIO:
+    # draw mini-arrows for field visualization
+    for i, ((vx, vy), lab) in enumerate(zip(vectors, labels)):
+        if i % 15 != 0:
             continue
-
-        # 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)
+        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)
 
     combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
     out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
@@ -113,32 +105,60 @@ def classify_wrong_direction(traj_json, flow_model_json, bg_img=None):
     return out_path
 
 
-# ------------------------------------------------------------
-# 🖥️ Gradio Interface
-# ------------------------------------------------------------
+# ============================================================
+# 🚀 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
+# ============================================================
 description_text = """
-### 🚦 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°.
+### 🧭 Dominant Flow Learning (Stage 2 — Cosine-Based Improved)
+Upload the **trajectories JSON** from Stage 1.  
+Optionally upload a background frame for overlay visualization.
 """
 
+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=classify_wrong_direction,
+    fn=process_json,
     inputs=[
-        gr.File(label="Trajectories JSON (Stage 1)"),
-        gr.File(label="Flow Model JSON (Stage 2)"),
+        gr.File(label="Upload trajectories JSON"),
         gr.File(label="Optional background frame (.jpg)")
     ],
-    outputs=gr.Image(label="Wrong-Direction Overlay"),
-    title="🚗 Wrong-Direction Detection – Stage 3 (Angle + Temporal + Zones)",
-    description=description_text
+    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,
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()