Angle-Based Flow Learning

app.py

@@ -19,48 +19,51 @@ def extract_motion_vectors(data):
 
 
 # ============================================================
-# 🧮 2. Dominant Flow Clustering (Cosine-based + Auto-Fix)
+# 🧮 2. Direction-Specific (Angle-Based) Clustering
 # ============================================================
-def learn_flows_improved(vectors, n_clusters=2):
-    """Cosine-based clustering of normalized motion directions."""
+def cluster_by_angle(vectors, n_clusters=2):
+    """Cluster motion directions using circular (angle-space) logic."""
     if len(vectors) < n_clusters:
         return None, None
 
-    # --- Normalize & filter weak motions ---
-    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) < n_clusters:
-        return None, None
+    # --- Convert to angles (−180° → 180°) ---
+    angles = np.degrees(np.arctan2(vectors[:, 1], vectors[:, 0]))
+    angles = angles.reshape(-1, 1)
 
-    # --- Cluster ---
+    # --- Run clustering in angle space ---
     kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=42)
-    kmeans.fit(dirs)
-    centers = kmeans.cluster_centers_
-    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
-
-    # --- ✅ Auto-fix: ensure at least one pair of flows are truly opposite ---
-    if len(centers) >= 2:
-        sim = np.dot(centers[0], centers[1])
-        if sim > -0.8:
-            # If not opposite enough, flip the smaller-magnitude cluster
-            if np.linalg.norm(centers[0]) < np.linalg.norm(centers[1]):
-                centers[0] = -centers[0]
-            else:
-                centers[1] = -centers[1]
-
-    # --- Assign labels ---
-    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
+    kmeans.fit(angles)
+    centers = kmeans.cluster_centers_.flatten()
+
+    # --- Convert centers back to unit direction vectors ---
+    centers_rad = np.radians(centers)
+    flow_vectors = np.column_stack((np.cos(centers_rad), np.sin(centers_rad)))
+
+    # --- Ensure flows are sufficiently opposite (auto-flip if needed) ---
+    if len(flow_vectors) >= 2:
+        sim = np.dot(flow_vectors[0], flow_vectors[1])
+        if sim > -0.8:  # not opposite enough
+            flow_vectors[1] = -flow_vectors[0]
+
+    # --- Assign labels based on closest angular distance ---
+    def angle_distance(a, b):
+        d = np.abs(a - b)
+        return np.minimum(d, 360 - d)
+
+    labels = np.zeros(len(angles), dtype=int)
+    for i, ang in enumerate(angles.flatten()):
+        d0 = angle_distance(ang, centers[0])
+        d1 = angle_distance(ang, centers[1]) if n_clusters > 1 else 999
+        labels[i] = 0 if d0 < d1 else 1
+
+    return labels, flow_vectors
 
 
 # ============================================================
 # 🧭 3. Estimate Road Angle from Dominant Flow
 # ============================================================
 def estimate_road_angle(centers):
-    """Return average flow direction in degrees (0° = horizontal right)."""
+    """Return average flow direction in degrees (0° = right)."""
     if centers is None or len(centers) == 0:
         return 0.0
     dominant = np.mean(centers, axis=0)
@@ -81,12 +84,11 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None,
         bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
 
     overlay = bg.copy()
-    colors = [(0, 0, 255), (255, 255, 0), (0, 255, 255), (255, 0, 255)]
+    colors = [(0, 0, 255), (255, 255, 0)]
 
     norms = np.linalg.norm(vectors, axis=1, keepdims=True)
     vectors = np.divide(vectors, norms + 1e-6) * 10
 
-    # --- Sampled arrows for visual flow density ---
     for i, ((vx, vy), lab) in enumerate(zip(vectors, labels)):
         if i % 15 != 0:
             continue
@@ -95,7 +97,6 @@ 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)
 
-    # --- Draw dominant flow arrows ---
     h, w = overlay.shape[:2]
     scale = 300
     center_pt = (w // 2, h // 2)
@@ -110,7 +111,6 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None,
         cv2.putText(overlay, f"Flow {i+1}", (end[0] + 10, end[1]),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
 
-    # --- Optional zones overlay ---
     if drive_zone is not None:
         cv2.polylines(overlay, [np.array(drive_zone, np.int32)], True, (0, 255, 255), 2)
         cv2.putText(overlay, "Drive Zone", tuple(np.array(drive_zone[0], int)),
@@ -140,17 +140,14 @@ def process_json(json_file, background=None):
     if len(vectors) == 0:
         return None, {"error": "No motion vectors found."}
 
-    # --- Use 2 clusters normally, can bump to 3 if road has multiple flows ---
-    labels, centers = learn_flows_improved(vectors, n_clusters=2)
+    labels, centers = cluster_by_angle(vectors, n_clusters=2)
     if labels is None:
         return None, {"error": "Insufficient data for clustering."}
 
     road_angle = estimate_road_angle(centers)
 
     drive_zone = [[100, 100], [800, 100], [800, 500], [100, 500]]
-    entry_zones = [
-        [[50, 100], [100, 100], [100, 500], [50, 500]]
-    ]
+    entry_zones = [[[50, 100], [100, 100], [100, 500], [50, 500]]]
 
     img_path = draw_flow_overlay(vectors, labels, centers,
                                  background, drive_zone, entry_zones)
@@ -170,9 +167,9 @@ def process_json(json_file, background=None):
 # 🖥️ 6. Gradio Interface
 # ============================================================
 description_text = """
-### 🧭 Dominant Flow Learning (Stage 2 — Angle + Zone-Aware + Auto-Fix)
-Uploads the **trajectories JSON** from Stage 1 and optionally a background frame.  
-Outputs dominant flow directions (auto-corrected if not opposite), estimated road angle, and zone polygons for Stage 3.
+### 🧭 Dominant Flow Learning (Stage 2 — Angle-Based)
+Clusters vehicle motion **by direction angle** on a circular scale,  
+giving cleaner opposite flows even on curved or diagonal roads.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -188,7 +185,7 @@ demo = gr.Interface(
         gr.Image(label="Dominant Flow Overlay"),
         gr.JSON(label="Flow Stats (Stage 2 Output)")
     ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Auto-Fix for Opposite Flows)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Angle-Based)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )