show_only_wrong

app.py

@@ -1,5 +1,5 @@
 # ============================================================
-# 🚦 Stage 3 — Wrong Direction Detection (Improved)
+# 🚦 Stage 3 — Wrong Direction Detection (Improved + Toggle)
 # ============================================================
 
 import os, cv2, json, tempfile, numpy as np, gradio as gr
@@ -44,7 +44,7 @@ class Track:
         return [x, y]
 
 # ============================================================
-# ⚙️ Utilities
+# ⚙️ Utility Functions
 # ============================================================
 def compute_cosine_similarity(v1, v2):
     v1 = v1 / (np.linalg.norm(v1) + 1e-6)
@@ -64,7 +64,7 @@ def smooth_direction(points, window=5):
 # ============================================================
 # 🧭 Wrong-Direction Detection Core
 # ============================================================
-def process_video(video_file, stage2_json):
+def process_video(video_file, stage2_json, show_only_wrong=False):
     data = json.load(open(stage2_json))
     lane_flows = np.array(data.get("flow_centers", [[1,0]]))
     drive_zone = np.array(data.get("drive_zone", []))
@@ -131,16 +131,19 @@ def process_video(video_file, stage2_json):
             sims = [compute_cosine_similarity(motion, f) for f in lane_flows]
             best_sim = max(sims)
 
-            # only classify after some frames (to reduce false early flag)
+            # only classify after some frames
             if trk.frames_seen > DELAY_FRAMES:
                 if best_sim < SIM_THRESH:
                     trk.status = "WRONG"
-                    color = (0, 0, 255)
                 else:
                     trk.status = "OK"
-                    color = (0, 255, 0)
-                cv2.putText(frame, f"ID:{tid} {trk.status}", tuple(np.int32(pos)),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+
+                # draw depending on toggle
+                if (not show_only_wrong) or (trk.status == "WRONG"):
+                    color = (0, 0, 255) if trk.status == "WRONG" else (0, 255, 0)
+                    cv2.putText(frame, f"ID:{tid} {trk.status}",
+                                tuple(np.int32(pos)),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
         out.write(frame)
 
@@ -153,16 +156,18 @@ def process_video(video_file, stage2_json):
 # ============================================================
 description = """
 ### 🚦 Stage 3 — Wrong Direction Detection (Improved)
-- Uses cosine similarity instead of raw angle comparison  
-- Lane-wise flow support for curved roads  
-- Temporal smoothing & delayed classification  
+- ✅ Cosine similarity instead of raw angle  
+- ✅ Lane-wise flow support for curved roads  
+- ✅ Temporal smoothing & delayed classification  
+- ✅ Toggle to show only WRONG vehicles  
 """
 
 demo = gr.Interface(
     fn=process_video,
     inputs=[
         gr.File(label="Input Video"),
-        gr.File(label="Stage 2 Flow JSON")
+        gr.File(label="Stage 2 Flow JSON"),
+        gr.Checkbox(label="Show ONLY Wrong Labels Overlay", value=False)
     ],
     outputs=gr.Video(label="Output (with WRONG/OK labels)"),
     title="🚗 Stage 3 – Improved Wrong-Direction Detection",