Revert

app.py

@@ -1,161 +1,211 @@
-# ============================================================
-# 🚦 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
+import os, cv2, json, tempfile, zipfile, numpy as np, gradio as gr
+from ultralytics import YOLO
+from filterpy.kalman import KalmanFilter
+from scipy.optimize import linear_sum_assignment
 
 # ------------------------------------------------------------
-# ⚙️ CONFIG
+# 🔧  Safe-load fix for PyTorch 2.6
 # ------------------------------------------------------------
-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
+import torch, ultralytics.nn.tasks as ultralytics_tasks
+torch.serialization.add_safe_globals([ultralytics_tasks.DetectionModel])
 
 # ------------------------------------------------------------
-# 🔧 Helper – universal loader for Gradio inputs
+# ⚙️  YOLO setup
 # ------------------------------------------------------------
-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.")
+MODEL_PATH = "yolov8n.pt"
+model = YOLO(MODEL_PATH)
+VEHICLE_CLASSES = [2, 3, 5, 7]  # car, motorcycle, bus, truck
+
 
 # ------------------------------------------------------------
-# 🧩 Load Stage 2 flow model
+# 🧩  Kalman tracker
 # ------------------------------------------------------------
-def load_flow_model(flow_model_json):
-    model = load_json_input(flow_model_json)
-    centers = [np.array(z) for z in model["zone_flow_centers"]]
-    return centers
+class Track:
+    def __init__(self, bbox, tid):
+        self.id = tid
+        self.kf = KalmanFilter(dim_x=4, dim_z=2)
+        self.kf.F = np.array([[1,0,1,0],[0,1,0,1],[0,0,1,0],[0,0,0,1]])
+        self.kf.H = np.array([[1,0,0,0],[0,1,0,0]])
+        self.kf.P *= 1000.0
+        self.kf.R *= 10.0
+        self.kf.x[:2] = np.array(self.centroid(bbox)).reshape(2,1)
+        self.trace = []
+
+    def centroid(self, b):
+        x1, y1, x2, y2 = b
+        return [(x1+x2)/2, (y1+y2)/2]
+
+    def predict(self):
+        self.kf.predict()
+        return self.kf.x[:2].reshape(2)
+
+    def update(self, b):
+        z = np.array(self.centroid(b)).reshape(2,1)
+        self.kf.update(z)
+        cx, cy = self.kf.x[:2].reshape(2)
+        self.trace.append((float(cx), float(cy)))
+        return (cx, cy)
+
 
 # ------------------------------------------------------------
-# 🧩 Extract trajectories (Stage 1)
+# 🧮  Direction analyzer
 # ------------------------------------------------------------
-def extract_trajectories(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
+def analyze_direction(trace, centers):
+    if len(trace) < 3:
+        return "NA", 1.0
+    v = np.array(trace[-1]) - np.array(trace[-3])
+    if np.linalg.norm(v) < 1e-6:
+        return "NA", 1.0
+    v = v / np.linalg.norm(v)
+    sims = np.dot(centers, v)
+    max_sim = np.max(sims)
+    if max_sim < 0:
+        return "WRONG", float(max_sim)
+    return "OK", float(max_sim)
+
 
 # ------------------------------------------------------------
-# 🧮 Direction + Angle + Confidence Helpers
+# 🧭  Load normalized flow centers
 # ------------------------------------------------------------
-def smooth_direction(pts, window=SMOOTH_FRAMES):
-    if len(pts) < 2: return np.array([0,0])
-    diffs = np.diff(pts[-window:], axis=0)
-    v = np.mean(diffs, axis=0)
-    return v / (np.linalg.norm(v)+1e-6)
-
-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)
-    return np.degrees(np.arccos(cosang))
-
-def angle_to_confidence(angle):
-    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))
+def load_flow_centers(flow_json):
+    data = json.load(open(flow_json))
+    centers = np.array(data["flow_centers"])
+    centers = centers / (np.linalg.norm(centers, axis=1, keepdims=True) + 1e-6)
+    return centers
+
 
 # ------------------------------------------------------------
-# 🎥 Main logic → annotated video
+# 🎥  Process video
 # ------------------------------------------------------------
-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)
-
-    # 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
-    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())
+def process_video(video_path, flow_json, show_only_wrong=False):
+    centers = load_flow_centers(flow_json)
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25
+    w, h = int(cap.get(3)), int(cap.get(4))
+
     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
-
-    # 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
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
+
+    tracks, next_id, log = [], 0, []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        results = model(frame, verbose=False)[0]
+        detections = []
+        for box in results.boxes:
+            if int(box.cls) in VEHICLE_CLASSES and box.conf > 0.3:
+                detections.append(box.xyxy[0].cpu().numpy())
+
+        # Predict existing
+        predicted = [t.predict() for t in tracks]
+        predicted = np.array(predicted) if len(predicted) > 0 else np.empty((0,2))
+
+        # Assign detections to tracks
+        assigned = set()
+        if len(predicted) > 0 and len(detections) > 0:
+            cost = np.zeros((len(predicted), len(detections)))
+            for i, p in enumerate(predicted):
+                for j, d in enumerate(detections):
+                    cx, cy = ((d[0]+d[2])/2, (d[1]+d[3])/2)
+                    cost[i,j] = np.linalg.norm(p - np.array([cx,cy]))
+            r, c = linear_sum_assignment(cost)
+            for i, j in zip(r, c):
+                if cost[i,j] < 80:
+                    assigned.add(j)
+                    tracks[i].update(detections[j])
+
+        # New tracks
+        for j, d in enumerate(detections):
+            if j not in assigned:
+                t = Track(d, next_id)
+                next_id += 1
+                t.update(d)
+                tracks.append(t)
+
+        # --- 🧩 Draw + Log (toggle support) ---
+        for trk in tracks:
+            if len(trk.trace) < 3:
+                continue
+            status, sim = analyze_direction(trk.trace, centers)
+
+            # Skip OKs if toggle is enabled
+            if show_only_wrong and status != "WRONG":
+                continue
+
+            x, y = map(int, trk.trace[-1])
+            color = (0,255,0) if status=="OK" else ((0,0,255) if status=="WRONG" else (200,200,200))
+            cv2.circle(frame,(x,y),4,color,-1)
+            cv2.putText(frame,f"ID:{trk.id} {status}",(x-20,y-10),
+                        cv2.FONT_HERSHEY_SIMPLEX,0.5,color,1)
+            for i in range(1,len(trk.trace)):
+                cv2.line(frame,
+                         (int(trk.trace[i-1][0]),int(trk.trace[i-1][1])),
+                         (int(trk.trace[i][0]),int(trk.trace[i][1])),
+                         color,1)
+
+            # Log once per unique vehicle
+            if len(trk.trace) > 5 and not any(entry["id"] == trk.id for entry in log):
+                log.append({"id": trk.id, "status": status, "cos_sim": round(sim,3)})
+
+        out.write(frame)
+
+    cap.release()
+    out.release()
+
+    # Unique summary
+    unique_ids = {entry["id"] for entry in log}
+    summary = {"vehicles_analyzed": len(unique_ids)}
+
+    # Create ZIP bundle
+    zip_path = tempfile.NamedTemporaryFile(suffix=".zip", delete=False).name
+    with zipfile.ZipFile(zip_path, "w") as zf:
+        zf.write(out_path, arcname="violation_output.mp4")
+        zf.writestr("per_vehicle_log.json", json.dumps(log, indent=2))
+        zf.writestr("summary.json", json.dumps(summary, indent=2))
+
+    return out_path, log, summary, zip_path
+
 
 # ------------------------------------------------------------
-# 🖥️ Gradio Interface
+# 🖥️  Gradio interface
 # ------------------------------------------------------------
+def run_app(video, flow_file, show_only_wrong):
+    vid, log_json, summary, zip_file = process_video(video, flow_file, show_only_wrong)
+    return vid, log_json, summary, zip_file
+
+
 description_text = """
-### 🚦 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
+### 🚦 Wrong-Direction Detection (Stage 3)
+Upload your traffic video and the **flow_stats.json** from Stage 2.  
+You can toggle whether to display all detections or only WRONG-direction vehicles.
 """
 
 demo = gr.Interface(
-    fn=classify_wrong_direction_video,
+    fn=run_app,
     inputs=[
-        gr.File(label="Trajectories JSON (Stage 1)"),
-        gr.File(label="Flow Model JSON (Stage 2)"),
-        gr.File(label="Optional background frame (.jpg/.png)")
+        gr.Video(label="Upload Traffic Video (.mp4)"),
+        gr.File(label="Upload flow_stats.json (Stage 2 Output)"),
+        gr.Checkbox(label="Show Only Wrong Labels", value=False)
+    ],
+    outputs=[
+        gr.Video(label="Violation Output Video"),
+        gr.JSON(label="Per-Vehicle Log"),
+        gr.JSON(label="Summary"),
+        gr.File(label="⬇️ Download All Outputs (ZIP)")
     ],
-    outputs=gr.Video(label="Annotated Video Output"),
-    title="🚗 Stage 3 — Wrong-Direction Detection (Video Output)",
-    description=description_text
+    title="🚗 Wrong-Direction Detection – Stage 3 (Toggle + ZIP)",
+    description=description_text,
+    examples=None,
 )
 
+# Disable analytics / flagging / SSR
+demo.flagging_mode = "never"
+demo.cache_examples = False
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False, show_api=False)