Update app.py

app.py

@@ -1,384 +1,914 @@
 """
-RF-DETR Truck Counter — counts trucks crossing a fixed horizontal line
-at the center of the video (either direction counts as one crossing).
+Bridge Traffic & Load Demo App
+Fast RF-DETR + ByteTrack vehicle counting for bridge videos.
 """
 
 import os
-import tempfile
 import time
-from collections import defaultdict
+import tempfile
+from functools import lru_cache
+from typing import Dict, List, Tuple
 
 import cv2
 import gradio as gr
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import supervision as sv
 import torch
-from rfdetr import RFDETRNano
-
-# ---------------------------------------------------------------------------
-# Target classes (COCO indices)
-# ---------------------------------------------------------------------------
-# We still DETECT all of these so users see them tracked on screen,
-# but only TRUCKS contribute to the line-crossing count.
-TARGET_CLASSES = {
-    0:  "person",
-    2:  "car",
-    7:  "truck",
-    16: "dog",
-    17: "horse",
-    18: "sheep",
-    19: "cow",
+
+from rfdetr import RFDETRNano, RFDETRMedium
+
+
+# ---------------------------------------------------------------------
+# Vehicle classes from COCO
+# ---------------------------------------------------------------------
+# COCO IDs used by RF-DETR:
+# 2 = car, 3 = motorcycle, 5 = bus, 7 = truck
+VEHICLE_CLASSES: Dict[int, str] = {
+    2: "car",
+    3: "motorcycle",
+    5: "bus",
+    7: "truck",
 }
-TARGET_IDS = list(TARGET_CLASSES.keys())
-TRUCK_CLASS_ID = 7
-
-DISPLAY_NAMES = {
-    "person": "person",
-    "car":    "car",
-    "truck":  "truck",
-    "dog":    "dog",
-    "horse":  "horse / donkey",
-    "sheep":  "sheep / goat",
-    "cow":    "cow",
+
+# Very rough demonstration weights in kg.
+# Adjust these for your local traffic profile.
+DEFAULT_WEIGHTS_KG: Dict[int, int] = {
+    2: 1500,    # car / small vehicle
+    3: 250,     # motorcycle
+    5: 12000,   # bus
+    7: 18000,   # truck / lorry
 }
 
-CLASS_COLORS = {  # BGR
-    "person":  (66, 135, 245),
-    "car":     (66, 245, 167),
-    "truck":   (245, 66, 161),
-    "dog":     (120, 245, 200),
-    "horse":   (245, 120, 120),
-    "sheep":   (220, 220, 220),
-    "cow":     (140, 90, 60),
+CLASS_COLORS_BGR: Dict[int, Tuple[int, int, int]] = {
+    2: (40, 190, 120),    # car
+    3: (255, 170, 70),    # motorcycle
+    5: (245, 120, 45),    # bus
+    7: (220, 70, 180),    # truck
 }
 
-APP_DIR = os.path.dirname(os.path.abspath(__file__))
-EXAMPLE_VIDEO = os.path.join(APP_DIR, "example.mp4")
+MODEL_OPTIONS = {
+    "Nano - fastest": RFDETRNano,
+    "Medium - more accurate, slower": RFDETRMedium,
+}
 
-# ---------------------------------------------------------------------------
-# Load model — CPU-pinned
-# ---------------------------------------------------------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Loading RF-DETR Nano on {DEVICE}…")
-MODEL = RFDETRNano(device=DEVICE)
-
-if DEVICE == "cuda":
-    try:
-        MODEL.optimize_for_inference()
-        print("Optimized for GPU inference.")
-    except Exception as e:
-        print(f"GPU optimization skipped: {e}")
 
 try:
     torch.set_num_threads(max(1, (os.cpu_count() or 2) - 1))
 except Exception:
     pass
 
-print("Model ready.")
-
+if DEVICE == "cuda":
+    torch.backends.cudnn.benchmark = True
 
-# ---------------------------------------------------------------------------
-# Draw a custom truck counter directly on the centerline
-# ---------------------------------------------------------------------------
-def draw_truck_label_on_line(frame: np.ndarray, line_y: int, total: int) -> np.ndarray:
-    text = f"TRUCKS CROSSED: {total}"
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    scale = max(0.6, frame.shape[1] / 1600)
-    thickness = max(2, int(2 * scale))
 
-    (tw, th), baseline = cv2.getTextSize(text, font, scale, thickness)
-    pad_x, pad_y = 14, 10
-    box_w = tw + 2 * pad_x
-    box_h = th + 2 * pad_y + baseline
+# ---------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------
+@lru_cache(maxsize=2)
+def load_model(model_name: str):
+    """Load RF-DETR once and reuse it across runs."""
+    model_cls = MODEL_OPTIONS[model_name]
 
-    cx = frame.shape[1] // 2
-    x1 = cx - box_w // 2
-    y1 = line_y - box_h // 2
-    x2 = x1 + box_w
-    y2 = y1 + box_h
+    print(f"Loading {model_name} on {DEVICE}...")
 
+    try:
+        model = model_cls(device=DEVICE)
+    except TypeError:
+        # Fallback for older RF-DETR builds
+        model = model_cls()
+
+    if DEVICE == "cuda":
+        try:
+            model.optimize_for_inference()
+            print("RF-DETR optimized for inference.")
+        except Exception as exc:
+            print(f"optimize_for_inference skipped: {exc}")
+
+    print("Model ready.")
+    return model
+
+
+# ---------------------------------------------------------------------
+# Detection helper
+# ---------------------------------------------------------------------
+def predict_vehicles(
+    model,
+    frame_bgr: np.ndarray,
+    confidence: float,
+    inference_width: int,
+) -> sv.Detections:
+    """
+    Resize frame before inference for speed, then scale boxes back to
+    original video coordinates.
+    """
+    h, w = frame_bgr.shape[:2]
+    inference_width = int(inference_width)
+
+    if inference_width > 0 and w > inference_width:
+        scale = inference_width / float(w)
+        resized_w = inference_width
+        resized_h = int(h * scale)
+        model_frame = cv2.resize(frame_bgr, (resized_w, resized_h), interpolation=cv2.INTER_AREA)
+    else:
+        scale = 1.0
+        model_frame = frame_bgr
+
+    frame_rgb = cv2.cvtColor(model_frame, cv2.COLOR_BGR2RGB)
+
+    with torch.inference_mode():
+        detections = model.predict(frame_rgb, threshold=float(confidence))
+
+    if len(detections) == 0:
+        return detections
+
+    # Keep only vehicle classes.
+    mask = np.isin(detections.class_id, list(VEHICLE_CLASSES.keys()))
+    detections = detections[mask]
+
+    if len(detections) == 0:
+        return detections
+
+    # Scale boxes back to original frame size.
+    if scale != 1.0:
+        detections.xyxy = detections.xyxy / scale
+
+    return detections
+
+
+# ---------------------------------------------------------------------
+# Counting and load helpers
+# ---------------------------------------------------------------------
+def side_of_line(y: float, line_y: int, dead_zone_px: int = 4) -> int:
+    """
+    Returns -1 above the line, +1 below the line, 0 inside a small dead zone.
+    The dead zone prevents jitter around the line from causing false crossings.
+    """
+    diff = y - line_y
+    if abs(diff) <= dead_zone_px:
+        return 0
+    return -1 if diff < 0 else 1
+
+
+def detection_centres(detections: sv.Detections) -> np.ndarray:
+    if len(detections) == 0:
+        return np.empty((0, 2), dtype=float)
+    xyxy = detections.xyxy
+    cx = (xyxy[:, 0] + xyxy[:, 2]) / 2.0
+    cy = (xyxy[:, 1] + xyxy[:, 3]) / 2.0
+    return np.column_stack([cx, cy])
+
+
+def get_class_weight_kg(class_id: int, weights: Dict[int, int]) -> int:
+    return int(weights.get(int(class_id), 0))
+
+
+def draw_header_panel(
+    frame: np.ndarray,
+    total_count: int,
+    cumulative_kg: float,
+    live_load_kg: float,
+    load_index_percent: float,
+    fps_text: str,
+) -> np.ndarray:
+    """Draw a clean dashboard panel at the top-left of the frame."""
     overlay = frame.copy()
-    cv2.rectangle(overlay, (x1, y1), (x2, y2), (245, 66, 161), -1)  # truck-pink
-    frame = cv2.addWeighted(overlay, 0.88, frame, 0.12, 0)
-    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 255), 2)
-
-    text_x = x1 + pad_x
-    text_y = y1 + pad_y + th
-    cv2.putText(frame, text, (text_x, text_y),
-                font, scale, (255, 255, 255), thickness, cv2.LINE_AA)
+    x1, y1, x2, y2 = 18, 18, 520, 158
+    cv2.rectangle(overlay, (x1, y1), (x2, y2), (20, 24, 36), -1)
+    frame = cv2.addWeighted(overlay, 0.82, frame, 0.18, 0)
+
+    cv2.putText(frame, "BRIDGE TRAFFIC LOAD DEMO", (34, 46),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.72, (255, 255, 255), 2, cv2.LINE_AA)
+
+    cv2.putText(frame, f"Vehicles crossed: {total_count}", (34, 78),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.62, (230, 240, 255), 2, cv2.LINE_AA)
+
+    cv2.putText(frame, f"Cumulative estimated mass: {cumulative_kg / 1000.0:.1f} tonnes", (34, 106),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.58, (220, 240, 230), 2, cv2.LINE_AA)
+
+    cv2.putText(frame, f"Live load: {live_load_kg / 1000.0:.1f} t | Load index: {load_index_percent:.1f}% | {fps_text}", (34, 134),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.52, (230, 230, 255), 1, cv2.LINE_AA)
+
     return frame
 
 
-# ---------------------------------------------------------------------------
-# HTML side panel
-# ---------------------------------------------------------------------------
-def build_counts_html(truck_total: int, truck_in: int, truck_out: int,
-                      frame_idx: int, total: int, elapsed: float) -> str:
-    pct = (frame_idx / total * 100) if total else 0
-
-    hero = (
-        '<div style="text-align:center;padding:22px 14px;margin-bottom:14px;'
-        'background:linear-gradient(135deg,#f5318b,#c2185b);color:white;'
-        'border-radius:14px;box-shadow:0 4px 12px rgba(245,49,139,0.35);">'
-        '<div style="font-size:32px;line-height:1;">🚛</div>'
-        '<div style="font-size:11px;opacity:0.9;letter-spacing:1.5px;'
-        'margin-top:6px;">TRUCKS CROSSED</div>'
-        f'<div style="font-size:54px;font-weight:800;line-height:1.0;'
-        f'margin-top:4px;">{truck_total}</div>'
-        '<div style="display:flex;justify-content:center;gap:14px;'
-        'margin-top:10px;font-size:12px;opacity:0.95;">'
-        f'<span>↓ {truck_in} down</span>'
-        f'<span>↑ {truck_out} up</span>'
-        '</div></div>'
+def annotate_frame(
+    frame: np.ndarray,
+    detections: sv.Detections,
+    line_y: int,
+    roi_top_y: int,
+    roi_bottom_y: int,
+    class_counts: Dict[str, int],
+    total_count: int,
+    cumulative_kg: float,
+    live_load_kg: float,
+    load_index_percent: float,
+    fps_text: str,
+) -> np.ndarray:
+    """Draw ROI, counting line, boxes, labels and dashboard."""
+    h, w = frame.shape[:2]
+
+    # Bridge deck ROI overlay
+    overlay = frame.copy()
+    cv2.rectangle(overlay, (0, roi_top_y), (w, roi_bottom_y), (80, 80, 80), -1)
+    frame = cv2.addWeighted(overlay, 0.08, frame, 0.92, 0)
+
+    # Counting line
+    cv2.line(frame, (0, line_y), (w, line_y), (40, 230, 255), 3)
+    cv2.putText(frame, "COUNTING LINE", (24, max(28, line_y - 12)),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.58, (40, 230, 255), 2, cv2.LINE_AA)
+
+    # ROI borders
+    cv2.line(frame, (0, roi_top_y), (w, roi_top_y), (170, 170, 170), 1)
+    cv2.line(frame, (0, roi_bottom_y), (w, roi_bottom_y), (170, 170, 170), 1)
+
+    if len(detections) > 0:
+        tracker_ids = detections.tracker_id
+        if tracker_ids is None:
+            tracker_ids = [None] * len(detections)
+
+        confidences = detections.confidence
+        if confidences is None:
+            confidences = [0.0] * len(detections)
+
+        for xyxy, class_id, conf, track_id in zip(
+            detections.xyxy,
+            detections.class_id,
+            confidences,
+            tracker_ids,
+        ):
+            class_id = int(class_id)
+            x1, y1, x2, y2 = map(int, xyxy)
+            name = VEHICLE_CLASSES.get(class_id, "vehicle")
+            color = CLASS_COLORS_BGR.get(class_id, (80, 220, 255))
+            weight_t = DEFAULT_WEIGHTS_KG.get(class_id, 0) / 1000.0
+
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+
+            id_text = f"#{int(track_id)} " if track_id is not None and int(track_id) >= 0 else ""
+            label = f"{id_text}{name} {float(conf):.2f} ~{weight_t:.1f}t"
+            (tw, th), base = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.52, 1)
+
+            label_y1 = max(0, y1 - th - base - 8)
+            cv2.rectangle(frame, (x1, label_y1), (x1 + tw + 10, y1), color, -1)
+            cv2.putText(frame, label, (x1 + 5, y1 - 6),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.52, (255, 255, 255), 1, cv2.LINE_AA)
+
+    frame = draw_header_panel(
+        frame=frame,
+        total_count=total_count,
+        cumulative_kg=cumulative_kg,
+        live_load_kg=live_load_kg,
+        load_index_percent=load_index_percent,
+        fps_text=fps_text,
     )
 
-    progress = (
-        '<div style="margin:8px 0 4px 0;">'
-        '<div style="display:flex;justify-content:space-between;font-size:11px;'
-        'color:#6b7280;margin-bottom:4px;">'
-        f'<span>frame {frame_idx} / {total}</span>'
-        f'<span>{pct:.1f}% · {elapsed:.1f}s</span>'
-        '</div>'
-        '<div style="height:6px;background:#e5e7eb;border-radius:3px;overflow:hidden;">'
-        f'<div style="height:100%;width:{pct}%;background:#6366f1;transition:width 0.2s;"></div>'
-        '</div></div>'
-    )
+    # Compact class counts at bottom
+    items = [f"{k}: {v}" for k, v in class_counts.items() if v > 0]
+    count_text = " | ".join(items) if items else "No crossings yet"
+    cv2.putText(frame, count_text, (22, h - 24),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.62, (255, 255, 255), 2, cv2.LINE_AA)
 
-    return hero + progress
+    return frame
 
 
-def build_summary_md(truck_total: int, truck_in: int, truck_out: int) -> str:
-    if truck_total == 0:
-        return ("### ℹ️ No trucks crossed the center line.\n"
-                "Try a lower confidence threshold or a smaller frame stride.")
-    return (f"### 🚛 Total trucks crossed: **{truck_total}**\n\n"
-            f"- ↓ Going down: {truck_in}\n"
-            f"- ↑ Going up: {truck_out}")
+def build_metrics_html(
+    total_count: int,
+    class_counts: Dict[str, int],
+    cumulative_kg: float,
+    live_load_kg: float,
+    load_index_percent: float,
+    frame_idx: int,
+    total_frames: int,
+    elapsed: float,
+    device: str,
+) -> str:
+    pct = (frame_idx / total_frames * 100.0) if total_frames else 0.0
+    tonnes = cumulative_kg / 1000.0
+    live_tonnes = live_load_kg / 1000.0
+
+    car = class_counts.get("car", 0)
+    motorcycle = class_counts.get("motorcycle", 0)
+    bus = class_counts.get("bus", 0)
+    truck = class_counts.get("truck", 0)
+
+    return f"""
+    <div style="font-family:Inter,system-ui,Arial;">
+      <div style="display:grid;grid-template-columns:1fr 1fr;gap:10px;margin-bottom:12px;">
+        <div style="padding:16px;border-radius:16px;background:linear-gradient(135deg,#1d4ed8,#312e81);color:white;">
+          <div style="font-size:11px;letter-spacing:1px;opacity:.85;">VEHICLES CROSSED</div>
+          <div style="font-size:46px;font-weight:800;line-height:1;">{total_count}</div>
+        </div>
+        <div style="padding:16px;border-radius:16px;background:linear-gradient(135deg,#be185d,#7e22ce);color:white;">
+          <div style="font-size:11px;letter-spacing:1px;opacity:.85;">EST. CUMULATIVE MASS</div>
+          <div style="font-size:36px;font-weight:800;line-height:1;">{tonnes:.1f} t</div>
+        </div>
+      </div>
+
+      <div style="display:grid;grid-template-columns:1fr 1fr;gap:10px;margin-bottom:12px;">
+        <div style="padding:14px;border:1px solid #e5e7eb;border-radius:14px;background:white;">
+          <div style="font-size:12px;color:#6b7280;">Live bridge load</div>
+          <div style="font-size:28px;font-weight:750;color:#111827;">{live_tonnes:.1f} t</div>
+        </div>
+        <div style="padding:14px;border:1px solid #e5e7eb;border-radius:14px;background:white;">
+          <div style="font-size:12px;color:#6b7280;">Load index</div>
+          <div style="font-size:28px;font-weight:750;color:#111827;">{load_index_percent:.1f}%</div>
+        </div>
+      </div>
+
+      <div style="padding:14px;border:1px solid #e5e7eb;border-radius:14px;background:#ffffff;margin-bottom:12px;">
+        <div style="font-size:12px;color:#6b7280;margin-bottom:8px;">Crossings by class</div>
+        <div style="display:grid;grid-template-columns:1fr 1fr;gap:8px;font-size:14px;">
+          <div>🚗 Cars: <b>{car}</b></div>
+          <div>🏍️ Motorcycles: <b>{motorcycle}</b></div>
+          <div>🚌 Buses: <b>{bus}</b></div>
+          <div>🚛 Trucks: <b>{truck}</b></div>
+        </div>
+      </div>
+
+      <div style="font-size:12px;color:#6b7280;margin-bottom:4px;display:flex;justify-content:space-between;">
+        <span>Frame {frame_idx} / {total_frames}</span>
+        <span>{pct:.1f}% · {elapsed:.1f}s · {device}</span>
+      </div>
+      <div style="height:8px;background:#e5e7eb;border-radius:99px;overflow:hidden;">
+        <div style="height:100%;width:{pct:.2f}%;background:#4f46e5;"></div>
+      </div>
+    </div>
+    """
+
+
+def render_load_plot(history: List[Dict]) -> np.ndarray:
+    """Render load-index chart as an RGB image for Gradio."""
+    if not history:
+        img = np.ones((320, 600, 3), dtype=np.uint8) * 255
+        cv2.putText(img, "Load index chart will appear here", (60, 165),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (80, 80, 80), 2, cv2.LINE_AA)
+        return img
+
+    df = pd.DataFrame(history)
+    # Plot only a manageable number of points for speed.
+    if len(df) > 500:
+        df = df.iloc[np.linspace(0, len(df) - 1, 500).astype(int)]
+
+    fig, ax = plt.subplots(figsize=(8.0, 3.8), dpi=100)
+    ax.plot(df["time_s"], df["load_index_percent"], linewidth=2)
+    ax.set_title("Estimated Bridge Load Index Over Time")
+    ax.set_xlabel("Video time (seconds)")
+    ax.set_ylabel("Load index (%)")
+    ax.grid(True, alpha=0.25)
+    ax.set_ylim(bottom=0)
+    fig.tight_layout()
+
+    fig.canvas.draw()
+    rgba = np.asarray(fig.canvas.buffer_rgba())
+    rgb = cv2.cvtColor(rgba, cv2.COLOR_RGBA2RGB)
+    plt.close(fig)
+    return rgb
+
+
+def build_final_summary(
+    total_count: int,
+    class_counts: Dict[str, int],
+    cumulative_kg: float,
+    peak_live_load_kg: float,
+    peak_load_index: float,
+    csv_path: str,
+) -> str:
+    tonnes = cumulative_kg / 1000.0
+    peak_tonnes = peak_live_load_kg / 1000.0
+
+    return f"""
+### Final bridge traffic summary
+
+**Vehicles crossed:** {total_count}
+
+| Vehicle class | Count |
+|---|---:|
+| Cars | {class_counts.get("car", 0)} |
+| Motorcycles | {class_counts.get("motorcycle", 0)} |
+| Buses | {class_counts.get("bus", 0)} |
+| Trucks | {class_counts.get("truck", 0)} |
+
+**Cumulative estimated mass:** {tonnes:.2f} tonnes  
+**Peak estimated live load:** {peak_tonnes:.2f} tonnes  
+**Peak load index:** {peak_load_index:.1f}%
+
+The CSV output contains the estimated load-index time series for later plotting or analysis.
+
+> Note: This is a demonstration traffic-load indicator, not a certified structural stress calculation.
+"""
 
 
-# ---------------------------------------------------------------------------
-# Main streaming generator
-# ---------------------------------------------------------------------------
-def process_video(video_path, confidence, frame_stride):
+# ---------------------------------------------------------------------
+# Main processing generator
+# ---------------------------------------------------------------------
+def process_video(
+    video_path,
+    model_name,
+    confidence,
+    frame_stride,
+    inference_width,
+    line_position_percent,
+    roi_top_percent,
+    roi_bottom_percent,
+    reference_capacity_tonnes,
+    car_weight_t,
+    motorcycle_weight_t,
+    bus_weight_t,
+    truck_weight_t,
+):
     if video_path is None:
-        yield (None,
-               '<div style="padding:12px;color:#b91c1c;">⚠️ Please upload a video first.</div>',
-               "Submit a video to start.",
-               None)
+        yield (
+            None,
+            build_metrics_html(0, {"car": 0, "motorcycle": 0, "bus": 0, "truck": 0}, 0, 0, 0, 0, 0, 0, DEVICE),
+            render_load_plot([]),
+            "Upload a video to start analysis.",
+            None,
+            None,
+        )
         return
 
-    video_info = sv.VideoInfo.from_video_path(video_path)
-    width, height = video_info.width, video_info.height
-
-    # ---- Line zone fixed at vertical center ----
-    line_y = height // 2
-    line_zone = sv.LineZone(
-        start=sv.Point(0, line_y),
-        end=sv.Point(width, line_y),
+    # Update demo weights from UI.
+    weights_kg = {
+        2: int(float(car_weight_t) * 1000),
+        3: int(float(motorcycle_weight_t) * 1000),
+        5: int(float(bus_weight_t) * 1000),
+        7: int(float(truck_weight_t) * 1000),
+    }
+    # Keep global-like drawing labels consistent for this run.
+    DEFAULT_WEIGHTS_KG.update(weights_kg)
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise RuntimeError(f"Could not open video: {video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps is None or fps <= 1:
+        fps = 25.0
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
+
+    if width <= 0 or height <= 0:
+        cap.release()
+        raise RuntimeError("Could not read video dimensions.")
+
+    line_y = int(height * float(line_position_percent) / 100.0)
+    roi_top_y = int(height * float(roi_top_percent) / 100.0)
+    roi_bottom_y = int(height * float(roi_bottom_percent) / 100.0)
+
+    if roi_bottom_y <= roi_top_y:
+        roi_top_y = int(height * 0.25)
+        roi_bottom_y = int(height * 0.90)
+
+    reference_capacity_kg = max(1.0, float(reference_capacity_tonnes) * 1000.0)
+
+    yield (
+        None,
+        build_metrics_html(0, {"car": 0, "motorcycle": 0, "bus": 0, "truck": 0}, 0, 0, 0, 0, total_frames, 0, DEVICE),
+        render_load_plot([]),
+        "### Loading RF-DETR model and starting analysis...",
+        None,
+        None,
     )
 
-    # ---- Annotators (sized to frame) ----
-    scale = max(0.5, width / 1280)
-    box_ann = sv.BoxAnnotator(thickness=max(2, int(2 * scale)))
-    label_ann = sv.LabelAnnotator(
-        text_scale=0.5 * scale,
-        text_thickness=max(1, int(1 * scale)),
-        text_padding=4,
+    model = load_model(str(model_name))
+    tracker = sv.ByteTrack(frame_rate=int(round(fps)))
+
+    # Output files
+    out_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    out_csv_path = tempfile.NamedTemporaryFile(suffix=".csv", delete=False).name
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(out_video_path, fourcc, fps, (width, height))
+
+    # State
+    last_detections = sv.Detections.empty()
+    last_side_by_id: Dict[int, int] = {}
+    counted_ids = set()
+    track_class: Dict[int, int] = {}
+
+    class_counts = {"car": 0, "motorcycle": 0, "bus": 0, "truck": 0}
+    total_count = 0
+    cumulative_kg = 0.0
+
+    history: List[Dict] = []
+    event_rows: List[Dict] = []
+
+    start_wall = time.time()
+    last_yield_wall = 0.0
+    last_plot = render_load_plot([])
+    processed_frames = 0
+
+    peak_live_load_kg = 0.0
+    peak_load_index = 0.0
+
+    frame_idx = 0
+
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+
+        detect_this_frame = (frame_idx % int(frame_stride) == 0)
+
+        if detect_this_frame:
+            detections = predict_vehicles(
+                model=model,
+                frame_bgr=frame,
+                confidence=float(confidence),
+                inference_width=int(inference_width),
+            )
+            detections = tracker.update_with_detections(detections)
+            last_detections = detections
+        else:
+            detections = last_detections
+
+        # Update per-track class and line crossing only when we have tracked detections.
+        centres = detection_centres(detections)
+
+        live_load_kg = 0.0
+        active_track_ids = set()
+
+        if len(detections) > 0 and detections.tracker_id is not None:
+            for det_i, (centre, class_id, track_id) in enumerate(
+                zip(centres, detections.class_id, detections.tracker_id)
+            ):
+                if track_id is None or int(track_id) < 0:
+                    continue
+
+                tid = int(track_id)
+                cid = int(class_id)
+                cy = float(centre[1])
+
+                track_class[tid] = cid
+                active_track_ids.add(tid)
+
+                # Live bridge-deck load, only if the object is inside the deck ROI.
+                if roi_top_y <= cy <= roi_bottom_y:
+                    live_load_kg += get_class_weight_kg(cid, weights_kg)
+
+                current_side = side_of_line(cy, line_y)
+                previous_side = last_side_by_id.get(tid)
+
+                if current_side != 0:
+                    if previous_side is not None and previous_side != 0:
+                        crossed = previous_side != current_side
+                        if crossed and tid not in counted_ids:
+                            vehicle_name = VEHICLE_CLASSES.get(cid, "vehicle")
+                            vehicle_weight = get_class_weight_kg(cid, weights_kg)
+                            direction = "down" if previous_side < current_side else "up"
+
+                            counted_ids.add(tid)
+                            total_count += 1
+                            class_counts[vehicle_name] = class_counts.get(vehicle_name, 0) + 1
+                            cumulative_kg += vehicle_weight
+
+                            event_rows.append({
+                                "video_time_s": frame_idx / fps,
+                                "frame": frame_idx,
+                                "tracker_id": tid,
+                                "vehicle_type": vehicle_name,
+                                "direction": direction,
+                                "estimated_vehicle_weight_kg": vehicle_weight,
+                                "cumulative_estimated_mass_kg": cumulative_kg,
+                            })
+
+                    last_side_by_id[tid] = current_side
+
+        load_index_percent = (live_load_kg / reference_capacity_kg) * 100.0
+        peak_live_load_kg = max(peak_live_load_kg, live_load_kg)
+        peak_load_index = max(peak_load_index, load_index_percent)
+
+        history.append({
+            "video_time_s": frame_idx / fps,
+            "time_s": frame_idx / fps,
+            "frame": frame_idx,
+            "vehicles_crossed_total": total_count,
+            "cars_crossed": class_counts.get("car", 0),
+            "motorcycles_crossed": class_counts.get("motorcycle", 0),
+            "buses_crossed": class_counts.get("bus", 0),
+            "trucks_crossed": class_counts.get("truck", 0),
+            "live_load_kg": live_load_kg,
+            "live_load_tonnes": live_load_kg / 1000.0,
+            "load_index_percent": load_index_percent,
+            "cumulative_estimated_mass_kg": cumulative_kg,
+            "cumulative_estimated_mass_tonnes": cumulative_kg / 1000.0,
+        })
+
+        elapsed_wall = time.time() - start_wall
+        processed_frames += 1
+        current_processing_fps = processed_frames / max(elapsed_wall, 1e-6)
+        fps_text = f"{current_processing_fps:.1f} proc FPS"
+
+        annotated = annotate_frame(
+            frame=frame,
+            detections=detections,
+            line_y=line_y,
+            roi_top_y=roi_top_y,
+            roi_bottom_y=roi_bottom_y,
+            class_counts=class_counts,
+            total_count=total_count,
+            cumulative_kg=cumulative_kg,
+            live_load_kg=live_load_kg,
+            load_index_percent=load_index_percent,
+            fps_text=fps_text,
+        )
+        writer.write(annotated)
+
+        now = time.time()
+        if now - last_yield_wall >= 0.35:
+            last_yield_wall = now
+            # Refresh the chart less often than the frame display.
+            last_plot = render_load_plot(history)
+            rgb_frame = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
+            yield (
+                rgb_frame,
+                build_metrics_html(
+                    total_count=total_count,
+                    class_counts=class_counts,
+                    cumulative_kg=cumulative_kg,
+                    live_load_kg=live_load_kg,
+                    load_index_percent=load_index_percent,
+                    frame_idx=frame_idx + 1,
+                    total_frames=total_frames,
+                    elapsed=elapsed_wall,
+                    device=DEVICE,
+                ),
+                last_plot,
+                "### Live analysis running...",
+                None,
+                None,
+            )
+
+        frame_idx += 1
+
+    cap.release()
+    writer.release()
+
+    # Save CSV time series. Add event-level detail as separate columns where possible.
+    history_df = pd.DataFrame(history)
+    history_df.to_csv(out_csv_path, index=False)
+
+    final_plot = render_load_plot(history)
+    final_summary = build_final_summary(
+        total_count=total_count,
+        class_counts=class_counts,
+        cumulative_kg=cumulative_kg,
+        peak_live_load_kg=peak_live_load_kg,
+        peak_load_index=peak_load_index,
+        csv_path=out_csv_path,
     )
-    trace_ann = sv.TraceAnnotator(thickness=max(2, int(2 * scale)), trace_length=40)
-    # We draw the line ourselves so we control the label completely
-    line_ann = sv.LineZoneAnnotator(
-        thickness=max(2, int(3 * scale)),
-        text_thickness=1, text_scale=0.01,    # effectively hide default text
-        display_in_count=False,
-        display_out_count=False,
+
+    final_frame = None
+    if history:
+        # Try to show the last annotated frame from the output video.
+        cap2 = cv2.VideoCapture(out_video_path)
+        if cap2.isOpened():
+            cap2.set(cv2.CAP_PROP_POS_FRAMES, max(0, frame_idx - 1))
+            ok, last = cap2.read()
+            if ok:
+                final_frame = cv2.cvtColor(last, cv2.COLOR_BGR2RGB)
+            cap2.release()
+
+    yield (
+        final_frame,
+        build_metrics_html(
+            total_count=total_count,
+            class_counts=class_counts,
+            cumulative_kg=cumulative_kg,
+            live_load_kg=0,
+            load_index_percent=0,
+            frame_idx=total_frames if total_frames else frame_idx,
+            total_frames=total_frames if total_frames else frame_idx,
+            elapsed=time.time() - start_wall,
+            device=DEVICE,
+        ),
+        final_plot,
+        final_summary,
+        out_video_path,
+        out_csv_path,
     )
 
-    frame_gen = sv.get_video_frames_generator(video_path)
-    tracker = sv.ByteTrack(frame_rate=int(video_info.fps))
 
-    last_detections = sv.Detections.empty()
-    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    total = video_info.total_frames or 0
-    start_time = time.time()
-    last_yield = 0.0
-    last_rgb = None
-
-    yield (None,
-           build_counts_html(0, 0, 0, 0, total, 0.0),
-           "### 🎬 Starting analysis…",
-           None)
-
-    with sv.VideoSink(target_path=out_path, video_info=video_info) as sink:
-        for i, frame in enumerate(frame_gen):
-
-            # ---- Detection (every Nth frame) ----
-            if i % frame_stride == 0:
-                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                detections = MODEL.predict(rgb, threshold=confidence)
-
-                if len(detections) > 0:
-                    mask = np.isin(detections.class_id, TARGET_IDS)
-                    detections = detections[mask]
-
-                detections = tracker.update_with_detections(detections)
-                last_detections = detections
-
-                # ---- Only trucks feed the line zone ----
-                if len(detections) > 0:
-                    truck_mask = detections.class_id == TRUCK_CLASS_ID
-                    truck_detections = detections[truck_mask]
-                    if len(truck_detections) > 0:
-                        line_zone.trigger(truck_detections)
-            else:
-                detections = last_detections
-
-            # ---- Annotate everything detected (visual richness) ----
-            if len(detections) > 0:
-                tids = (detections.tracker_id
-                        if detections.tracker_id is not None
-                        else [None] * len(detections))
-                labels = []
-                for cid, tid, conf in zip(detections.class_id, tids,
-                                          detections.confidence):
-                    name = TARGET_CLASSES.get(int(cid), "obj")
-                    display = DISPLAY_NAMES.get(name, name)
-                    tid_str = f"#{tid} " if tid is not None else ""
-                    labels.append(f"{tid_str}{display} {conf:.2f}")
-
-                frame = trace_ann.annotate(scene=frame, detections=detections)
-                frame = box_ann.annotate(scene=frame, detections=detections)
-                frame = label_ann.annotate(scene=frame, detections=detections,
-                                            labels=labels)
-
-            # ---- Draw line + custom truck counter on the line ----
-            frame = line_ann.annotate(frame=frame, line_counter=line_zone)
-            truck_in = int(line_zone.in_count)
-            truck_out = int(line_zone.out_count)
-            truck_total = truck_in + truck_out
-            frame = draw_truck_label_on_line(frame, line_y, truck_total)
-
-            sink.write_frame(frame)
-            last_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-
-            # ---- Yield to UI (throttled ~5/sec) ----
-            now = time.time()
-            if now - last_yield > 0.20 or i == total - 1:
-                last_yield = now
-                elapsed = time.time() - start_time
-                yield (last_rgb,
-                       build_counts_html(truck_total, truck_in, truck_out,
-                                         i + 1, total, elapsed),
-                       "### 🔴 Live analysis in progress…",
-                       None)
-
-    elapsed = time.time() - start_time
-    truck_in = int(line_zone.in_count)
-    truck_out = int(line_zone.out_count)
-    truck_total = truck_in + truck_out
-    yield (last_rgb,
-           build_counts_html(truck_total, truck_in, truck_out, total, total, elapsed),
-           build_summary_md(truck_total, truck_in, truck_out),
-           out_path)
-
-
-# ---------------------------------------------------------------------------
-# UI
-# ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------
 CUSTOM_CSS = """
-.gradio-container {max-width: 1240px !important; margin: auto;}
-#title-row {text-align: center; padding: 8px 0 0 0;}
-#title-row h1 {font-weight: 700; letter-spacing: -0.5px; margin-bottom: 4px;}
-#title-row p {color: #6b7280; margin-top: 0;}
-.card {border: 1px solid #e5e7eb; border-radius: 14px; padding: 16px;
-       background: #ffffff;}
-#live-frame img {border-radius: 10px;}
-footer {visibility: hidden;}
+.gradio-container {
+    max-width: 1320px !important;
+    margin: auto !important;
+}
+#hero {
+    text-align: center;
+    padding: 18px 8px 8px 8px;
+}
+#hero h1 {
+    font-weight: 850;
+    letter-spacing: -0.6px;
+    margin-bottom: 2px;
+}
+#hero p {
+    color: #64748b;
+    font-size: 16px;
+    margin-top: 0;
+}
+.panel {
+    border: 1px solid #e5e7eb;
+    border-radius: 18px;
+    padding: 16px;
+    background: #ffffff;
+    box-shadow: 0 8px 24px rgba(15, 23, 42, 0.04);
+}
+#live-frame img, #load-plot img {
+    border-radius: 14px;
+}
+footer {
+    visibility: hidden;
+}
 """
 
-with gr.Blocks(title="RF-DETR Truck Counter") as demo:
+with gr.Blocks(
+    title="Bridge Traffic Load Demo",
+    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
+    css=CUSTOM_CSS,
+) as demo:
 
-    with gr.Row(elem_id="title-row"):
+    with gr.Row(elem_id="hero"):
         gr.Markdown(
             """
-            # 🚛 RF-DETR Truck Counter
-            Counts trucks crossing a horizontal line at the **center of the video** —
-            in either direction. Powered by
-            [RF-DETR Nano](https://github.com/roboflow/rf-detr) + ByteTrack +
-            `sv.LineZone`.
+            # 🌉 Bridge Traffic Load Demo
+            Fast RF-DETR vehicle detection, ByteTrack tracking, line-crossing counts,
+            estimated cumulative vehicle mass, and live bridge load-index over time.
             """
         )
 
     with gr.Row():
-        # ---------- Left: input ----------
         with gr.Column(scale=1):
-            with gr.Group(elem_classes="card"):
-                gr.Markdown("### 📥 Input")
+            with gr.Group(elem_classes="panel"):
+                gr.Markdown("### 1) Upload video")
                 video_input = gr.Video(
-                    label="Upload a video",
+                    label="Bridge traffic video",
                     sources=["upload"],
                     format="mp4",
                     height=260,
                 )
 
-                with gr.Accordion("⚙️ Advanced settings", open=False):
-                    confidence = gr.Slider(
-                        minimum=0.1, maximum=0.9, value=0.45, step=0.05,
-                        label="Confidence threshold",
+                start_btn = gr.Button("▶ Start analysis", variant="primary", size="lg")
+
+                gr.Markdown("### 2) Speed settings")
+                model_name = gr.Radio(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    value="Nano - fastest",
+                    label="RF-DETR model",
+                )
+                confidence = gr.Slider(
+                    minimum=0.10,
+                    maximum=0.90,
+                    value=0.40,
+                    step=0.05,
+                    label="Confidence threshold",
+                )
+                frame_stride = gr.Slider(
+                    minimum=1,
+                    maximum=10,
+                    value=3,
+                    step=1,
+                    label="Frame stride",
+                    info="Detect every Nth frame. 1 is most accurate. 3-5 is much faster.",
+                )
+                inference_width = gr.Slider(
+                    minimum=384,
+                    maximum=1280,
+                    value=640,
+                    step=64,
+                    label="Inference width",
+                    info="Lower is faster. Try 512 or 640 for CPU demos.",
+                )
+
+                with gr.Accordion("Bridge settings", open=False):
+                    line_position_percent = gr.Slider(
+                        minimum=10,
+                        maximum=90,
+                        value=55,
+                        step=1,
+                        label="Counting line vertical position (%)",
                     )
-                    frame_stride = gr.Slider(
-                        minimum=1, maximum=15, value=3, step=1,
-                        label="Frame stride (CPU speed)",
-                        info="Detect every Nth frame. Higher = faster.",
+                    roi_top_percent = gr.Slider(
+                        minimum=0,
+                        maximum=90,
+                        value=20,
+                        step=1,
+                        label="Bridge deck ROI top (%)",
+                    )
+                    roi_bottom_percent = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=90,
+                        step=1,
+                        label="Bridge deck ROI bottom (%)",
+                    )
+                    reference_capacity_tonnes = gr.Slider(
+                        minimum=5,
+                        maximum=200,
+                        value=40,
+                        step=5,
+                        label="Reference live-load capacity for demo index (tonnes)",
                     )
 
-                submit_btn = gr.Button("▶️ Start Live Analysis",
-                                       variant="primary", size="lg")
-
-                gr.Markdown("#### 🎬 Example video")
-                gr.Examples(
-                    examples=[[EXAMPLE_VIDEO]],
-                    inputs=video_input,
-                    examples_per_page=4,
+                with gr.Accordion("Estimated class weights", open=False):
+                    car_weight_t = gr.Number(value=1.5, label="Car weight estimate (tonnes)")
+                    motorcycle_weight_t = gr.Number(value=0.25, label="Motorcycle weight estimate (tonnes)")
+                    bus_weight_t = gr.Number(value=12.0, label="Bus weight estimate (tonnes)")
+                    truck_weight_t = gr.Number(value=18.0, label="Truck weight estimate (tonnes)")
+
+                gr.Markdown(
+                    """
+                    **For speed:** use **Nano**, inference width **512-640**, and frame stride **3-5**.
+                    Use **Medium** only when you need better detection and have a GPU.
+                    """
                 )
 
-        # ---------- Right: live view ----------
         with gr.Column(scale=2):
-            with gr.Group(elem_classes="card"):
-                with gr.Row():
-                    with gr.Column(scale=3):
-                        gr.Markdown("### 🔴 Live View")
-                        live_frame = gr.Image(
-                            show_label=False,
-                            elem_id="live-frame",
-                            height=440,
+            with gr.Group(elem_classes="panel"):
+                gr.Markdown("### Live annotated video")
+                live_frame = gr.Image(
+                    show_label=False,
+                    elem_id="live-frame",
+                    height=470,
+                )
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel"):
+                        gr.Markdown("### Live metrics")
+                        metrics_html = gr.HTML(
+                            value=build_metrics_html(
+                                total_count=0,
+                                class_counts={"car": 0, "motorcycle": 0, "bus": 0, "truck": 0},
+                                cumulative_kg=0,
+                                live_load_kg=0,
+                                load_index_percent=0,
+                                frame_idx=0,
+                                total_frames=0,
+                                elapsed=0,
+                                device=DEVICE,
+                            )
                         )
-                    with gr.Column(scale=1, min_width=240):
-                        gr.Markdown("### 📊 Live Counts")
-                        live_counts = gr.HTML(
-                            value=build_counts_html(0, 0, 0, 0, 0, 0)
+
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel"):
+                        gr.Markdown("### Load index over time")
+                        load_plot = gr.Image(
+                            show_label=False,
+                            elem_id="load-plot",
+                            height=310,
+                            value=render_load_plot([]),
                         )
 
-    # ---------- Bottom: final results ----------
     with gr.Row():
         with gr.Column(scale=1):
-            with gr.Group(elem_classes="card"):
-                gr.Markdown("### 📤 Final annotated video")
-                video_output = gr.Video(label="Download / replay", height=260)
+            with gr.Group(elem_classes="panel"):
+                gr.Markdown("### Final annotated video")
+                video_output = gr.Video(label="Replay / download annotated video", height=270)
         with gr.Column(scale=1):
-            with gr.Group(elem_classes="card"):
-                gr.Markdown("### 📈 Final summary")
-                summary_output = gr.Markdown("Run an analysis to see results.")
+            with gr.Group(elem_classes="panel"):
+                gr.Markdown("### Final summary")
+                summary_output = gr.Markdown("Run an analysis to see the final summary.")
+                csv_output = gr.File(label="Download load-index CSV")
 
-    submit_btn.click(
+    start_btn.click(
         fn=process_video,
-        inputs=[video_input, confidence, frame_stride],
-        outputs=[live_frame, live_counts, summary_output, video_output],
+        inputs=[
+            video_input,
+            model_name,
+            confidence,
+            frame_stride,
+            inference_width,
+            line_position_percent,
+            roi_top_percent,
+            roi_bottom_percent,
+            reference_capacity_tonnes,
+            car_weight_t,
+            motorcycle_weight_t,
+            bus_weight_t,
+            truck_weight_t,
+        ],
+        outputs=[
+            live_frame,
+            metrics_html,
+            load_plot,
+            summary_output,
+            video_output,
+            csv_output,
+        ],
     )
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=4).launch(
-        theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
-        css=CUSTOM_CSS,
-    )
+    demo.queue(max_size=3).launch()