revert

app.py

@@ -13,67 +13,54 @@ 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
+            if np.linalg.norm(d) > 1:   # ignore jitter / static points
                 vectors.append(d)
     return np.array(vectors)
 
 
 # ============================================================
-# 🧮 2. Direction-Specific (Angle-Based) Clustering
+# 🧮 2. Improved Dominant Flow Clustering (Cosine-based)
 # ============================================================
-def cluster_by_angle(vectors, n_clusters=2):
-    """Cluster motion directions using circular (angle-space) logic."""
+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
+    """
     if len(vectors) < n_clusters:
         return None, None
 
-    # Convert to angles (−180° → 180°)
-    angles = np.degrees(np.arctan2(vectors[:, 1], vectors[:, 0])).reshape(-1, 1)
-
-    kmeans = KMeans(n_clusters=n_clusters, n_init=20, random_state=42)
-    kmeans.fit(angles)
-    centers = kmeans.cluster_centers_.flatten()
-
-    # 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 opposite (auto-flip if needed)
-    if len(flow_vectors) >= 2:
-        sim = np.dot(flow_vectors[0], flow_vectors[1])
-        if sim > -0.8:
-            flow_vectors[1] = -flow_vectors[0]
+    # (1) Normalize to direction only
+    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+    dirs = vectors / (norms + 1e-6)
 
-    # Assign labels by smallest angular distance
-    def angle_distance(a, b):
-        d = np.abs(a - b)
-        return np.minimum(d, 360 - d)
+    # (2) Filter out tiny motions
+    valid = (norms[:, 0] > 1.5)
+    dirs = dirs[valid]
+    if len(dirs) < n_clusters:
+        return None, None
 
-    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
+    # (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_
 
-    return labels, flow_vectors
+    # (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)
 
-# ============================================================
-# 🧭 3. Estimate Road Angle from Dominant Flow
-# ============================================================
-def estimate_road_angle(centers):
-    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)
+    return labels, centers
 
 
 # ============================================================
-# 🎨 4. Visualization Utility
+# 🎨 3. Visualization Utility (Option A — Scaled-up Arrows)
 # ============================================================
-def draw_flow_overlay(vectors, labels, centers, bg_img=None,
-                      drive_zone=None, entry_zones=None):
-    # Load background or fallback canvas
+def draw_flow_overlay(vectors, labels, centers, bg_img=None):
+    # background
     if bg_img and os.path.exists(bg_img):
         bg = cv2.imread(bg_img)
         if bg is None:
@@ -81,23 +68,27 @@ 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)]
+    colors = [(0, 0, 255), (255, 255, 0)]   # red & yellow
 
-    # Draw sample motion vectors
+    # 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
-        start = (np.random.randint(0, w), np.random.randint(0, h))
+        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 % len(colors)], 1, tipLength=0.3)
+        cv2.arrowedLine(overlay, start, end, colors[lab % 2], 1, tipLength=0.3)
 
-    # Draw dominant flow arrows
+    # --- main dominant arrows ---
+    h, w = overlay.shape[:2]
     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),
@@ -108,17 +99,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)
 
-    # 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)),
-                    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)
@@ -126,7 +106,7 @@ def draw_flow_overlay(vectors, labels, centers, bg_img=None,
 
 
 # ============================================================
-# 🚀 5. Combined Pipeline (With Dynamic Zones)
+# 🚀 4. Combined Pipeline
 # ============================================================
 def process_json(json_file, background=None):
     try:
@@ -138,47 +118,28 @@ def process_json(json_file, background=None):
     if len(vectors) == 0:
         return None, {"error": "No motion vectors found."}
 
-    labels, centers = cluster_by_angle(vectors, n_clusters=2)
+    labels, centers = learn_flows_improved(vectors)
     if labels is None:
         return None, {"error": "Insufficient data for clustering."}
 
-    road_angle = estimate_road_angle(centers)
-
-    # --- 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)
+    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
+    img_path = draw_flow_overlay(vectors, labels, centers, background)
 
     stats = {
         "num_vectors": int(len(vectors)),
         "dominant_flows": int(len(centers)),
-        "flow_centers": centers.tolist(),
-        "road_angle_deg": road_angle,
-        "drive_zone": drive_zone,
-        "entry_zones": entry_zones
+        "flow_centers": centers.tolist()
     }
     return img_path, stats
 
 
 # ============================================================
-# 🖥️ 6. Gradio Interface
+# 🖥️ 5. Gradio Interface
 # ============================================================
 description_text = """
-### 🧭 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.
+### 🧭 Dominant Flow Learning (Stage 2 — Cosine-Based Improved)
+Upload the **trajectories JSON** from Stage 1.  
+Optionally upload a background frame for overlay visualization.
 """
 
 example_json = "trajectories_sample.json" if os.path.exists("trajectories_sample.json") else None
@@ -192,12 +153,12 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(label="Dominant Flow Overlay"),
-        gr.JSON(label="Flow Stats (Stage 2 Output)")
+        gr.JSON(label="Flow Stats")
     ],
-    title="🚗 Dominant Flow Learning – Stage 2 (Dynamic Zones)",
+    title="🚗 Dominant Flow Learning – Stage 2 (Cosine-Based Improved)",
     description=description_text,
     examples=[[example_json, example_bg]] if example_json else None,
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()