Angle-Based + Dynamic Zones

app.py

@@ -26,26 +26,24 @@ def cluster_by_angle(vectors, n_clusters=2):
     if len(vectors) < n_clusters:
         return None, None
 
-    # --- Convert to angles (−180° → 180°) ---
-    angles = np.degrees(np.arctan2(vectors[:, 1], vectors[:, 0]))
-    angles = angles.reshape(-1, 1)
+    # Convert to angles (−180° → 180°)
+    angles = np.degrees(np.arctan2(vectors[:, 1], vectors[:, 0])).reshape(-1, 1)
 
-    # --- Run clustering in angle space ---
     kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=42)
     kmeans.fit(angles)
     centers = kmeans.cluster_centers_.flatten()
 
-    # --- Convert centers back to unit direction vectors ---
+    # Convert cluster centers → unit 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) ---
+    # Ensure flows are 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
+        if sim > -0.8:
             flow_vectors[1] = -flow_vectors[0]
 
-    # --- Assign labels based on closest angular distance ---
+    # Assign labels by smallest angular distance
     def angle_distance(a, b):
         d = np.abs(a - b)
         return np.minimum(d, 360 - d)
@@ -63,7 +61,6 @@ def cluster_by_angle(vectors, n_clusters=2):
 # 🧭 3. Estimate Road Angle from Dominant Flow
 # ============================================================
 def estimate_road_angle(centers):
-    """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)
@@ -76,6 +73,7 @@ def estimate_road_angle(centers):
 # ============================================================
 def draw_flow_overlay(vectors, labels, centers, bg_img=None,
                       drive_zone=None, entry_zones=None):
+    # Load background or fallback canvas
     if bg_img and os.path.exists(bg_img):
         bg = cv2.imread(bg_img)
         if bg is None:
@@ -83,24 +81,23 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None,
     else:
         bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
 
+    h, w = bg.shape[:2]
     overlay = bg.copy()
     colors = [(0, 0, 255), (255, 255, 0)]
 
+    # Draw sample motion vectors
     norms = np.linalg.norm(vectors, axis=1, keepdims=True)
     vectors = np.divide(vectors, norms + 1e-6) * 10
-
     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]))
+        start = (np.random.randint(0, w), np.random.randint(0, h))
         end = (int(start[0] + vx), int(start[1] + vy))
         cv2.arrowedLine(overlay, start, end, colors[lab % len(colors)], 1, tipLength=0.3)
 
-    h, w = overlay.shape[:2]
+    # Draw dominant flow arrows
     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),
@@ -111,6 +108,7 @@ 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)
 
+    # Draw zones if provided
     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)),
@@ -128,7 +126,7 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None,
 
 
 # ============================================================
-# 🚀 5. Combined Pipeline
+# 🚀 5. Combined Pipeline (With Dynamic Zones)
 # ============================================================
 def process_json(json_file, background=None):
     try:
@@ -146,8 +144,19 @@ def process_json(json_file, background=None):
 
     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]]]
+    # --- determine frame size for zones ---
+    if background and os.path.exists(background):
+        bg_img = cv2.imread(background)
+        h, w = bg_img.shape[:2]
+    else:
+        # fallback for unknown resolution
+        w, h = 1280, 720
+
+    # --- dynamic zones based on frame width ---
+    drive_zone = [[100, 100], [w - 100, 100], [w - 100, 500], [100, 500]]
+    entry_zones = [
+        [[w - 100, 100], [w - 50, 100], [w - 50, 500], [w - 100, 500]]  # right-edge entry
+    ]
 
     img_path = draw_flow_overlay(vectors, labels, centers,
                                  background, drive_zone, entry_zones)
@@ -167,9 +176,9 @@ def process_json(json_file, background=None):
 # 🖥️ 6. Gradio Interface
 # ============================================================
 description_text = """
-### 🧭 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.
+### 🧭 Dominant Flow Learning (Stage 2 — Angle-Based + Dynamic Zones)
+Clusters vehicle motion **by direction angle** on a circular scale  
+and automatically places the Drive Zone and Entry Zone using frame width.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -185,7 +194,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-Based)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Dynamic Zones)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )