Stage 3 (Video Output Version)

app.py

@@ -1,165 +1,159 @@
 # ============================================================
-# 🚦 Stage 3 – Wrong-Direction Detection
-#     (Angle + Temporal + Zone-Aware + Entry Gating + Confidence)
+# 🚦 Stage 3 – Wrong-Direction Detection (Video Output Version)
 # ============================================================
 
 import gradio as gr
 import numpy as np, cv2, json, os, tempfile
 from collections import defaultdict
+import math
 
 # ------------------------------------------------------------
 # ⚙️ CONFIG
 # ------------------------------------------------------------
-ANGLE_THRESHOLD = 60        # degrees → above this = WRONG
-SMOOTH_FRAMES = 5           # frames for temporal smoothing
-ENTRY_ZONE_RATIO = 0.15     # top 15% = entry region (skip)
+ANGLE_THRESHOLD = 60        # deg → above = WRONG
+SMOOTH_FRAMES   = 5         # temporal smoothing
+ENTRY_ZONE_RATIO = 0.15     # skip top 15 %
 CONF_MIN, CONF_MAX = 0, 100
+FPS = 25                    # output video fps
 
 # ------------------------------------------------------------
-# 1️⃣ Load flow model (Stage 2)
+# 🔧 Helper – universal loader for Gradio inputs
+# ------------------------------------------------------------
+def load_json_input(file_obj):
+    if file_obj is None:
+        raise ValueError("No file provided.")
+    if isinstance(file_obj, dict) and "name" in file_obj:
+        path = file_obj["name"]
+        return json.load(open(path))
+    elif hasattr(file_obj, "name"):
+        return json.load(open(file_obj.name))
+    elif isinstance(file_obj, str):
+        return json.load(open(file_obj))
+    else:
+        raise ValueError("Unsupported file input type.")
+
+# ------------------------------------------------------------
+# 🧩 Load Stage 2 flow model
 # ------------------------------------------------------------
 def load_flow_model(flow_model_json):
-    model = json.load(open(flow_model_json))
+    model = load_json_input(flow_model_json)
     centers = [np.array(z) for z in model["zone_flow_centers"]]
     return centers
 
 # ------------------------------------------------------------
-# 2️⃣ Extract trajectories
+# 🧩 Extract trajectories (Stage 1)
 # ------------------------------------------------------------
 def extract_trajectories(json_file):
-    data = json.load(open(json_file))
+    data = load_json_input(json_file)
     tracks = {tid: np.array(pts) for tid, pts in data.items() if len(pts) > 2}
     return tracks
 
 # ------------------------------------------------------------
-# 3️⃣ Smoothed direction for a trajectory
+# 🧮 Direction + Angle + Confidence Helpers
 # ------------------------------------------------------------
 def smooth_direction(pts, window=SMOOTH_FRAMES):
-    if len(pts) < 2:
-        return np.array([0, 0])
+    if len(pts) < 2: return np.array([0,0])
     diffs = np.diff(pts[-window:], axis=0)
     v = np.mean(diffs, axis=0)
-    n = np.linalg.norm(v)
-    return v / (n + 1e-6)
+    return v / (np.linalg.norm(v)+1e-6)
 
-# ------------------------------------------------------------
-# 4️⃣ Compute angular difference (deg)
-# ------------------------------------------------------------
 def angle_between(v1, v2):
-    v1 = v1 / (np.linalg.norm(v1) + 1e-6)
-    v2 = v2 / (np.linalg.norm(v2) + 1e-6)
-    cosang = np.clip(np.dot(v1, v2), -1, 1)
+    v1 = v1 / (np.linalg.norm(v1)+1e-6)
+    v2 = v2 / (np.linalg.norm(v2)+1e-6)
+    cosang = np.clip(np.dot(v1,v2), -1,1)
     return np.degrees(np.arccos(cosang))
 
-# ------------------------------------------------------------
-# 5️⃣ Determine zone index for y
-# ------------------------------------------------------------
-def get_zone_idx(y, frame_h, n_zones):
-    zone_height = frame_h / n_zones
-    return int(np.clip(y // zone_height, 0, n_zones - 1))
-
-# ------------------------------------------------------------
-# 6️⃣ Confidence mapping
-# ------------------------------------------------------------
 def angle_to_confidence(angle):
-    """
-    0° → 100% confidence
-    ANGLE_THRESHOLD° → 50%
-    180° → 0%
-    """
-    if angle < 0:
-        return CONF_MIN
-    if angle >= 180:
-        return CONF_MIN
-    # linear mapping: smaller angle = higher confidence
-    conf = max(CONF_MIN, CONF_MAX - (angle / 180) * 100)
-    return round(conf, 1)
+    if angle<0: return CONF_MIN
+    if angle>=180: return CONF_MIN
+    conf = max(CONF_MIN, CONF_MAX - (angle/180)*100)
+    return round(conf,1)
+
+def get_zone_idx(y, frame_h, n_zones):
+    zone_h = frame_h/n_zones
+    return int(np.clip(y//zone_h, 0, n_zones-1))
 
 # ------------------------------------------------------------
-# 7️⃣ Main logic
+# 🎥 Main logic → annotated video
 # ------------------------------------------------------------
-def classify_wrong_direction(traj_json, flow_model_json, bg_img=None):
+def classify_wrong_direction_video(traj_json, flow_model_json, bg_img=None):
     tracks = extract_trajectories(traj_json)
     centers_by_zone = load_flow_model(flow_model_json)
 
-    if bg_img and os.path.exists(bg_img):
-        bg = cv2.imread(bg_img)
+    # background size
+    if bg_img:
+        if isinstance(bg_img, dict) and "name" in bg_img:
+            bg_path = bg_img["name"]
+        elif hasattr(bg_img,"name"):
+            bg_path = bg_img.name
+        else:
+            bg_path = bg_img
+        bg = cv2.imread(bg_path)
     else:
-        bg = np.ones((600, 900, 3), dtype=np.uint8) * 40
-    h, w = bg.shape[:2]
-
-    overlay = bg.copy()
+        bg = np.ones((600,900,3),dtype=np.uint8)*40
+    if bg is None: bg = np.ones((600,900,3),dtype=np.uint8)*40
+    h,w = bg.shape[:2]
+
+    # infer video length from longest track
+    max_len = max(len(p) for p in tracks.values())
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    writer = cv2.VideoWriter(out_path, fourcc, FPS, (w,h))
     font = cv2.FONT_HERSHEY_SIMPLEX
-    results = []
-
-    for tid, pts in tracks.items():
-        if len(pts) < 3:
-            continue
-        cur_pt = pts[-1]
-        y = cur_pt[1]
-        zone_idx = get_zone_idx(y, h, len(centers_by_zone))
-
-        # Skip entry region
-        if y < h * ENTRY_ZONE_RATIO:
-            continue
-
-        v = smooth_direction(pts)
-        centers = centers_by_zone[zone_idx]
-        angles = [angle_between(v, c) for c in centers]
-        best_angle = min(angles)
-
-        # Confidence & label
-        conf = angle_to_confidence(best_angle)
-        label = "OK" if best_angle < ANGLE_THRESHOLD else "WRONG"
-        color = (0, 255, 0) if label == "OK" else (0, 0, 255)
-
-        # Draw trajectory & label
-        for p1, p2 in zip(pts[:-1], pts[1:]):
-            cv2.line(overlay, tuple(p1.astype(int)), tuple(p2.astype(int)), color, 2)
-        cv2.circle(overlay, tuple(cur_pt.astype(int)), 5, color, -1)
-        cv2.putText(
-            overlay,
-            f"ID:{tid} {label} ({conf}%)",
-            (int(cur_pt[0]) + 5, int(cur_pt[1]) - 5),
-            font, 0.6, color, 2
-        )
-
-        results.append({
-            "id": tid,
-            "zone": int(zone_idx),
-            "angle": round(best_angle, 1),
-            "confidence": conf,
-            "label": label
-        })
-
-    combined = cv2.addWeighted(bg, 0.6, overlay, 0.4, 0)
-    out_path = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False).name
-    cv2.imwrite(out_path, combined)
-    return out_path, results
+
+    # render frame-by-frame
+    for fi in range(max_len):
+        frame = bg.copy()
+        for tid, pts in tracks.items():
+            if fi >= len(pts): continue
+            cur_pt = pts[fi]
+            y = cur_pt[1]
+            zone_idx = get_zone_idx(y, h, len(centers_by_zone))
+            if y < h*ENTRY_ZONE_RATIO: continue
+
+            # smooth direction using past window
+            win = pts[max(0,fi-SMOOTH_FRAMES):fi+1]
+            v = smooth_direction(win)
+            centers = centers_by_zone[zone_idx]
+            angles = [angle_between(v,c) for c in centers]
+            best_angle = min(angles)
+            conf = angle_to_confidence(best_angle)
+            label = "OK" if best_angle < ANGLE_THRESHOLD else "WRONG"
+            color = (0,255,0) if label=="OK" else (0,0,255)
+
+            # draw trajectory so far
+            for p1,p2 in zip(pts[:fi], pts[1:fi+1]):
+                cv2.line(frame, tuple(p1.astype(int)), tuple(p2.astype(int)), color, 2)
+            cv2.circle(frame, tuple(cur_pt.astype(int)), 5, color, -1)
+            cv2.putText(frame, f"ID:{tid} {label} ({conf}%)",
+                        (int(cur_pt[0])+5, int(cur_pt[1])-5),
+                        font, 0.55, color, 2)
+
+        writer.write(frame)
+
+    writer.release()
+    return out_path
 
 # ------------------------------------------------------------
 # 🖥️ Gradio Interface
 # ------------------------------------------------------------
 description_text = """
-### 🚦 Wrong-Direction Detection (Stage 3 — with Confidence)
-- Compares each vehicle’s motion to its zone’s dominant flow.  
-- Uses angular difference → smaller angle ⇒ higher confidence.  
-- Ignores entry region to avoid false positives.  
-- Displays ID, label, and confidence percentage.
+### 🚦 Stage 3 – Wrong-Direction Detection (Video Output)
+Uses **trajectories (Stage 1)** + **flow model (Stage 2)** to create an annotated MP4:  
+- Angle-based + temporal smoothing + zone awareness  
+- Entry-zone gating  
+- Confidence (%) per vehicle
 """
 
 demo = gr.Interface(
-    fn=classify_wrong_direction,
+    fn=classify_wrong_direction_video,
     inputs=[
         gr.File(label="Trajectories JSON (Stage 1)"),
         gr.File(label="Flow Model JSON (Stage 2)"),
-        gr.File(label="Optional background frame (.jpg)")
-    ],
-    outputs=[
-        gr.Image(label="Annotated Output"),
-        gr.JSON(label="Per-Vehicle Results")
+        gr.File(label="Optional background frame (.jpg/.png)")
     ],
-    title="🚗 Stage 3 — Wrong-Direction Detection (with Confidence)",
+    outputs=gr.Video(label="Annotated Video Output"),
+    title="🚗 Stage 3 — Wrong-Direction Detection (Video Output)",
     description=description_text
 )