Auto lane detection (#12)

- Auto lane detection (813043571749baabd651b27e6c2c8834fd75bccf)

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np, cv2, json, tempfile, os
 from sklearn.cluster import KMeans
 
+
 # ============================================================
 # 🧩 1. Compute motion vectors from trajectory JSON
 # ============================================================
@@ -13,54 +14,60 @@ def extract_motion_vectors(data):
             continue
         diffs = np.diff(pts, axis=0)
         for d in diffs:
-            if np.linalg.norm(d) > 1:   # ignore jitter / static points
+            if np.linalg.norm(d) > 1:   # ignore jitter / static
                 vectors.append(d)
     return np.array(vectors)
 
 
 # ============================================================
-# 🧮 2. Auto-Adaptive Dominant Flow Clustering (Cosine-based)
+# 🧮 2. Auto-Adaptive Dominant Flow Clustering (with dominance sort)
 # ============================================================
 def learn_flows_auto(vectors, normalize=True, max_clusters=2):
     """
-    Automatically determines whether there's 1 or 2 dominant flows:
-    - Computes angular spread of all motion vectors.
-    - If spread < 15°, uses 1 cluster (single-lane one-way).
-    - Otherwise, uses 2 clusters (two-way or opposing flows).
+    Automatically chooses 1 or 2 flow clusters depending on angular spread.
+    Also returns cluster membership counts to rank dominance.
     """
     if len(vectors) < 3:
-        return None, None
+        return None, None, None
 
-    # (1) Normalize to unit direction
+    # Normalize
     norms = np.linalg.norm(vectors, axis=1, keepdims=True)
     dirs = vectors / (norms + 1e-6)
     valid = (norms[:, 0] > 1.5)
     dirs = dirs[valid]
     if len(dirs) < 3:
-        return None, None
+        return None, None, None
 
-    # (2) Compute angular spread in degrees
+    # Angular spread check
     angles = np.degrees(np.arctan2(dirs[:, 1], dirs[:, 0]))
     angles = (angles + 360) % 360
-    spread = np.ptp(angles)  # peak-to-peak range
+    spread = np.ptp(angles)
     n_clusters = 1 if spread < 15 else max_clusters
 
-    # (3) Run KMeans with chosen cluster count
+    # Cluster
     kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=42)
     kmeans.fit(dirs)
+    labels = kmeans.labels_
     centers = kmeans.cluster_centers_
 
+    # Normalize centers again
     centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
-    sims = np.dot(dirs, centers.T)
-    labels = np.argmax(sims, axis=1)
 
-    return labels, centers
+    # --- Dominance sorting by cluster size ---
+    counts = np.bincount(labels)
+    order = np.argsort(-counts)  # descending
+    centers = centers[order]
+    # remap labels to new dominance order
+    remap = {old: new for new, old in enumerate(order)}
+    labels = np.array([remap[l] for l in labels])
+
+    return labels, centers, counts[order]
 
 
 # ============================================================
-# 🎨 3. Visualization Utility (Scaled-up Arrows)
+# 🎨 3. Visualization Utility
 # ============================================================
-def draw_flow_overlay(vectors, labels, centers, bg_img=None):
+def draw_flow_overlay(vectors, labels, centers, counts, bg_img=None):
     if bg_img and os.path.exists(bg_img):
         bg = cv2.imread(bg_img)
         if bg is None:
@@ -71,9 +78,11 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None):
     overlay = bg.copy()
     colors = [(0, 0, 255), (255, 255, 0), (255, 0, 255)]
 
+    # normalize arrow lengths
     norms = np.linalg.norm(vectors, axis=1, keepdims=True)
     vectors = np.divide(vectors, norms + 1e-6) * 10
 
+    # small random field arrows
     for i, ((vx, vy), lab) in enumerate(zip(vectors, labels)):
         if i % 15 != 0:
             continue
@@ -82,20 +91,27 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None):
         end = (int(start[0] + vx), int(start[1] + vy))
         cv2.arrowedLine(overlay, start, end, colors[lab % len(colors)], 1, tipLength=0.3)
 
-    # --- main dominant arrows ---
+    # --- Main dominant arrows ---
     h, w = overlay.shape[:2]
     scale = 300
     center_pt = (w // 2, h // 2)
 
-    for i, c in enumerate(centers):
+    for i, (c, count) in enumerate(zip(centers, counts)):
         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)
+        cv2.putText(
+            overlay,
+            f"Flow {i+1} ({count} vecs)",
+            (end[0] + 10, end[1]),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.7,
+            (0, 255, 0),
+            2,
+        )
 
     combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
     out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
@@ -116,17 +132,18 @@ def process_json(json_file, background=None):
     if len(vectors) == 0:
         return None, {"error": "No motion vectors found."}
 
-    labels, centers = learn_flows_auto(vectors)
+    labels, centers, counts = learn_flows_auto(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)
+    img_path = draw_flow_overlay(vectors, labels, centers, counts, background)
 
     stats = {
         "num_vectors": int(len(vectors)),
         "dominant_flows": int(len(centers)),
-        "flow_centers": centers.tolist()
+        "flow_counts": counts.tolist(),
+        "flow_centers": centers.tolist(),
     }
     return img_path, stats
 
@@ -135,10 +152,9 @@ def process_json(json_file, background=None):
 # 🖥️ 5. Gradio Interface
 # ============================================================
 description_text = """
-### 🧭 Dominant Flow Learning (Stage 2 — Auto-Adaptive)
-Automatically decides if traffic flow is **one-way** or **two-way**
-based on motion-vector angular spread.  
-Upload the **Stage 1 trajectories JSON**, and optionally a background frame.
+### 🧭 Dominant Flow Learning (Stage 2 – Auto + Dominance)
+Automatically detects if traffic is **one-way** or **two-way**  
+and orders flows by **vehicle count** so Flow 1 is the true dominant direction.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -154,7 +170,7 @@ demo = gr.Interface(
         gr.Image(label="Dominant Flow Overlay"),
         gr.JSON(label="Flow Stats")
     ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Auto-Adaptive)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Auto + Dominance)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )