Angle estimation (#6)

- Angle estimation (985c62c076c5799ace45d138fde3b9b897cf1679)

app.py

@@ -13,54 +13,53 @@ 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. Improved Dominant Flow Clustering (Cosine-based)
+# 🧮 2. 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
-    """
+def learn_flows_improved(vectors, n_clusters=2):
+    """Cosine-based clustering of normalized motion directions."""
     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)
+# 🧭 3. Estimate Road Angle from Dominant Flow
+# ============================================================
+def estimate_road_angle(centers):
+    """Return average flow direction in degrees (0° = horizontal right)."""
+    if centers is None or len(centers) == 0:
+        return 0.0
+    dominant = np.mean(centers, axis=0)
+    angle = np.degrees(np.arctan2(dominant[1], dominant[0]))
+    return float(angle % 360)
+
+
 # ============================================================
-def draw_flow_overlay(vectors, labels, centers, bg_img=None):
-    # background
+# 🎨 4. Visualization Utility
+# ============================================================
+def draw_flow_overlay(vectors, labels, centers, bg_img=None,
+                      drive_zone=None, entry_zones=None):
     if bg_img and os.path.exists(bg_img):
         bg = cv2.imread(bg_img)
         if bg is None:
@@ -69,13 +68,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)]   # red & yellow
+    colors = [(0, 0, 255), (255, 255, 0)]
 
-    # normalize arrow lengths for small samples
     norms = np.linalg.norm(vectors, axis=1, keepdims=True)
     vectors = np.divide(vectors, norms + 1e-6) * 10
 
-    # draw mini-arrows for field visualization
     for i, ((vx, vy), lab) in enumerate(zip(vectors, labels)):
         if i % 15 != 0:
             continue
@@ -84,7 +81,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 % 2], 1, tipLength=0.3)
 
-    # --- main dominant arrows ---
     h, w = overlay.shape[:2]
     scale = 300
     center_pt = (w // 2, h // 2)
@@ -99,6 +95,17 @@ 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)),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
+    if entry_zones:
+        for ez in entry_zones:
+            cv2.polylines(overlay, [np.array(ez, np.int32)], True, (0, 0, 255), 2)
+            cv2.putText(overlay, "Entry Gate", tuple(np.array(ez[0], int)),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
+
     combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
     out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
     cv2.imwrite(out_path, combined)
@@ -106,7 +113,7 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None):
 
 
 # ============================================================
-# 🚀 4. Combined Pipeline
+# 🚀 5. Combined Pipeline
 # ============================================================
 def process_json(json_file, background=None):
     try:
@@ -122,24 +129,35 @@ def process_json(json_file, background=None):
     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)
+    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
+    ]
+
+    img_path = draw_flow_overlay(vectors, labels, centers,
+                                 background, drive_zone, entry_zones)
 
     stats = {
         "num_vectors": int(len(vectors)),
         "dominant_flows": int(len(centers)),
-        "flow_centers": centers.tolist()
+        "flow_centers": centers.tolist(),
+        "road_angle_deg": road_angle,
+        "drive_zone": drive_zone,
+        "entry_zones": entry_zones
     }
     return img_path, stats
 
 
 # ============================================================
-# 🖥️ 5. Gradio Interface
+# 🖥️ 6. 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.
+### 🧭 Dominant Flow Learning (Stage 2 — Angle + Zone-Aware)
+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.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -153,12 +171,12 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(label="Dominant Flow Overlay"),
-        gr.JSON(label="Flow Stats")
+        gr.JSON(label="Flow Stats (Stage 2 Output)")
     ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Cosine-Based Improved)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Angle + Zone-Aware)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()