auto-fix check for opposite directions

app.py

@@ -19,13 +19,14 @@ def extract_motion_vectors(data):
 
 
 # ============================================================
-# 🧮 2. Dominant Flow Clustering (Cosine-based)
+# 🧮 2. Dominant Flow Clustering (Cosine-based + Auto-Fix)
 # ============================================================
 def learn_flows_improved(vectors, n_clusters=2):
     """Cosine-based clustering of normalized motion directions."""
     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)
@@ -33,11 +34,23 @@ def learn_flows_improved(vectors, n_clusters=2):
     if len(dirs) < n_clusters:
         return None, None
 
+    # --- Cluster ---
     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
@@ -68,19 +81,21 @@ 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)]
+    colors = [(0, 0, 255), (255, 255, 0), (0, 255, 255), (255, 0, 255)]
 
     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
         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)
+        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)
@@ -125,16 +140,16 @@ def process_json(json_file, background=None):
     if len(vectors) == 0:
         return None, {"error": "No motion vectors found."}
 
-    labels, centers = learn_flows_improved(vectors)
+    # --- Use 2 clusters normally, can bump to 3 if road has multiple flows ---
+    labels, centers = learn_flows_improved(vectors, n_clusters=2)
     if labels is None:
         return None, {"error": "Insufficient data for clustering."}
 
     road_angle = estimate_road_angle(centers)
 
-    # Optionally define default polygons (can be user-drawn later)
     drive_zone = [[100, 100], [800, 100], [800, 500], [100, 500]]
     entry_zones = [
-        [[50, 100], [100, 100], [100, 500], [50, 500]]  # left edge example
+        [[50, 100], [100, 100], [100, 500], [50, 500]]
     ]
 
     img_path = draw_flow_overlay(vectors, labels, centers,
@@ -155,9 +170,9 @@ def process_json(json_file, background=None):
 # 🖥️ 6. Gradio Interface
 # ============================================================
 description_text = """
-### 🧭 Dominant Flow Learning (Stage 2 — Angle + Zone-Aware)
+### 🧭 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, estimated road angle, and zone polygons for Stage 3.
+Outputs dominant flow directions (auto-corrected if not opposite), estimated road angle, and zone polygons for Stage 3.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -173,7 +188,7 @@ demo = gr.Interface(
         gr.Image(label="Dominant Flow Overlay"),
         gr.JSON(label="Flow Stats (Stage 2 Output)")
     ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Angle + Zone-Aware)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Auto-Fix for Opposite Flows)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )