app.py

app.py

@@ -0,0 +1,193 @@
+import gradio as gr
+import cv2, os, numpy as np, tempfile, time, json
+from ultralytics import YOLO
+from filterpy.kalman import KalmanFilter
+from scipy.optimize import linear_sum_assignment
+from tqdm import tqdm
+
+# ---------------------------------------------------------
+# ⚙️ INIT
+# ---------------------------------------------------------
+MODEL_PATH = "yolov8n.pt"
+model = YOLO(MODEL_PATH)
+
+# Vehicle classes from COCO
+VEHICLE_CLASSES = [2, 3, 5, 7]  # car, motorcycle, bus, truck
+
+
+# ---------------------------------------------------------
+# 🔍 SIMPLE KALMAN TRACKER
+# ---------------------------------------------------------
+class Track:
+    def __init__(self, bbox, track_id):
+        self.id = track_id
+        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.get_centroid(bbox)).reshape(2,1)
+        self.trace = []
+
+    def get_centroid(self,bbox):
+        x1,y1,x2,y2 = bbox
+        return [(x1+x2)/2,(y1+y2)/2]
+
+    def predict(self):
+        self.kf.predict()
+        return self.kf.x[:2].reshape(2)
+
+    def update(self,bbox):
+        z = np.array(self.get_centroid(bbox)).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)
+
+
+# ---------------------------------------------------------
+# 🎥 MAIN PROCESSOR
+# ---------------------------------------------------------
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    temp_out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    out = cv2.VideoWriter(temp_out.name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+
+    tracks = []
+    next_id = 0
+    trajectories = {}
+    frame_count = 0
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    pbar = tqdm(total=total_frames if total_frames>0 else 100, desc="Processing")
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_count += 1
+
+        # --- YOLO DETECTION ---
+        results = model(frame, verbose=False)[0]
+        detections = []
+        for box in results.boxes:
+            cls = int(box.cls)
+            if cls in VEHICLE_CLASSES and box.conf > 0.3:
+                detections.append(box.xyxy[0].cpu().numpy())
+
+        # --- PREDICT EXISTING TRACKS ---
+        predicted = [trk.predict() for trk in tracks]
+        predicted = np.array(predicted) if predicted else np.empty((0,2))
+
+        # --- ASSIGN DETECTIONS ---
+        assigned = set()
+        if len(predicted) > 0 and len(detections) > 0:
+            cost = np.zeros((len(predicted), len(detections)))
+            for i, trk in enumerate(predicted):
+                for j, det in enumerate(detections):
+                    cx, cy = ((det[0]+det[2])/2, (det[1]+det[3])/2)
+                    cost[i, j] = np.linalg.norm(trk - np.array([cx, cy]))
+            row_ind, col_ind = linear_sum_assignment(cost)
+            for r, c in zip(row_ind, col_ind):
+                if cost[r, c] < 80:  # distance threshold
+                    assigned.add(c)
+                    tracks[r].update(detections[c])
+
+        # --- NEW TRACKS ---
+        for j, det in enumerate(detections):
+            if j not in assigned:
+                trk = Track(det, next_id)
+                next_id += 1
+                trk.update(det)
+                tracks.append(trk)
+
+        # --- DRAW OUTPUT ---
+        for trk in tracks:
+            if len(trk.trace) < 2: 
+                continue
+            x,y = map(int,trk.trace[-1])
+            cv2.circle(frame,(x,y),3,(0,255,0),-1)
+            cv2.putText(frame,f"ID:{trk.id}",(x-10,y-10),cv2.FONT_HERSHEY_SIMPLEX,0.4,(0,255,0),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])),
+                                (0,255,0),1)
+            trajectories[trk.id] = trk.trace
+
+        out.write(frame)
+        pbar.update(1)
+
+    cap.release()
+    out.release()
+    pbar.close()
+
+    # Save trajectories JSON
+    traj_json = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
+    with open(traj_json.name, "w") as f:
+        json.dump(trajectories, f)
+
+    return temp_out.name, traj_json.name
+
+
+# ---------------------------------------------------------
+# 📤 WRAPPER FOR GRADIO
+# ---------------------------------------------------------
+def run_app(video_file):
+    # Copy uploaded video to temp path
+    temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+    if isinstance(video_file, dict) and "name" in video_file:
+        src_path = video_file["name"]
+    else:
+        src_path = video_file
+    with open(src_path, "rb") as src, open(temp_path, "wb") as dst:
+        dst.write(src.read())
+
+    start = time.time()
+    out_path, json_path = process_video(temp_path)
+    end = time.time()
+
+    summary = {
+        "total_time_sec": round(end-start,1),
+        "num_tracks": len(json.load(open(json_path))),
+        "avg_fps": round(cv2.VideoCapture(temp_path).get(cv2.CAP_PROP_FPS),2)
+    }
+
+    return out_path, json.load(open(json_path)), summary
+
+
+# ---------------------------------------------------------
+# 🖥️ GRADIO INTERFACE
+# ---------------------------------------------------------
+description_text = """
+### 🚦 Dominant Flow Tracker (Stage 1)
+Upload or select a sample traffic video below.  
+This app detects & tracks vehicles using YOLOv8 + Kalman Filter, and outputs:
+- Annotated tracking video
+- JSON trajectories
+- Summary stats for dominant-flow analysis
+"""
+
+example_video = "assets/examples/sample1.mp4" if os.path.exists("assets/examples/sample1.mp4") else None
+
+demo = gr.Interface(
+    fn=run_app,
+    inputs=gr.Video(label="Upload or use sample video (.mp4)", type="filepath"),
+    outputs=[
+        gr.Video(label="Tracked Output"),
+        gr.JSON(label="Vehicle Trajectories (Preview)"),
+        gr.JSON(label="Summary Stats")
+    ],
+    title="🚗 Dominant Flow Tracker – Stage 1",
+    description=description_text,
+    examples=[[example_video]] if example_video else None,
+)
+
+if __name__ == "__main__":
+    demo.launch()