Upload 2 files

lane_detection.py

@@ -0,0 +1,282 @@
+import cv2
+import numpy as np
+import time
+from ultralytics import YOLO
+import torch
+from collections import defaultdict
+import pandas as pd
+from typing import Dict, List, Tuple
+
+# LABEL_MAP = {
+#     0: "auto",
+#     1: "bus",
+#     2: "car",
+#     3: "motorcycle",
+#     4: "mini-bus",
+#     5: "scooter",
+#     6: "truck",
+# }
+LABEL_MAP = {0: "auto", 1: "bus", 2: "car", 3: "electric-rickshaw", 4: "large-sized-truck",5:'medium-sized-truck',6:'motorbike',7:'small-sized-truck'}
+
+def draw_text_with_background(
+    image,
+    text,
+    position,
+    font=cv2.FONT_HERSHEY_SIMPLEX,
+    font_scale=1,
+    font_thickness=2,
+    text_color=(255, 255, 255),
+    bg_color=(0, 0, 0),
+    padding=5,
+):
+    """Draw `text` on `image` with a filled rectangle behind it."""
+    (text_width, text_height), baseline = cv2.getTextSize(
+        text, font, font_scale, font_thickness
+    )
+    x, y = position
+    rect_y1 = y - text_height - padding - baseline // 2
+    rect_y2 = y + padding - baseline // 2
+
+    cv2.rectangle(
+        image,
+        (x, rect_y1),
+        (x + text_width + 2 * padding, rect_y2),
+        bg_color,
+        -1,
+    )
+    cv2.putText(
+        image,
+        text,
+        (x + padding, y - baseline // 2),
+        font,
+        font_scale,
+        text_color,
+        font_thickness,
+        cv2.LINE_AA,
+    )
+
+
+def get_color_for_class(cls_id: int):
+    """Deterministic bright color for each class index."""
+    np.random.seed(cls_id + 37)
+    return tuple(np.random.randint(100, 256, size=3).tolist())
+
+
+def _inside(pt: Tuple[int, int], poly: np.ndarray) -> bool:
+    """Point-in-polygon test using OpenCV (non‑zero if inside)."""
+    return cv2.pointPolygonTest(poly, pt, False) >= 0
+
+
+class YOLOVideoDetector:
+    """
+    Detect objects on a video and count them **per region**.
+
+    * `regions`: Dict[int, List[Tuple[int,int]]], mapping region id (0,1, …) to
+      4+ vertices in *pixel* coordinates (clockwise or anticlockwise).
+    * For each frame, counts are stored in a DataFrame column named
+      `<label>_<region>` (e.g. `car_0`, `bus_1`).
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        video_path: str,
+        output_path: str,
+        regions: Dict[int, List[Tuple[int, int]]],
+        classes=None,
+        conf: float = 0.35,
+        scale_factor: float = 1.5,
+    ):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+
+        self.model = YOLO(model_path)
+        self.video_path = video_path
+        self.output_path = output_path
+        self.conf = conf
+        self.classes = classes
+        self.scale = scale_factor
+
+        # ──────── NEW ────────
+        self.regions = {
+            rid: np.array(pts, np.int32) for rid, pts in regions.items() if pts
+        }
+        if not self.regions:
+            raise ValueError("`regions` cannot be empty — provide at least one polygon.")
+
+        # Prepare DataFrame columns once
+        self.df_columns = [
+            "Frame Number",
+            *[
+                f"{LABEL_MAP[c]}_{rid}"
+                for rid in self.regions
+                for c in LABEL_MAP.keys()
+            ],
+        ]
+
+    # ────────────────────────────────────────────────────────────────
+    def process_video(self) -> pd.DataFrame:
+        cap = cv2.VideoCapture(self.video_path)
+        if not cap.isOpened():
+            raise ValueError(f"Cannot open video: {self.video_path}")
+
+        ok, first_frame_original = cap.read()
+        if not ok:
+            cap.release()
+            raise ValueError(f"Cannot read first frame from: {self.video_path}")
+
+        h_orig, w_orig = first_frame_original.shape[:2]
+        prediction_counter_df = pd.DataFrame(columns=self.df_columns)
+
+        first_frame_processed = first_frame_original
+        frame_was_rotated = False
+
+        if w_orig < h_orig:
+            print(
+                f"Original frame (h,w): ({h_orig}, {w_orig}). Portrait → rotating 90° CW."
+            )
+            first_frame_processed = cv2.rotate(
+                first_frame_original, cv2.ROTATE_90_CLOCKWISE
+            )
+            frame_was_rotated = True
+        else:
+            print(f"Original frame (h,w): ({h_orig}, {w_orig}). Processing as landscape.")
+
+        # ----------------------------------------------------------------
+        base_h, base_w = first_frame_processed.shape[:2]
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+
+        out_w, out_h = int(base_w * self.scale), int(base_h * self.scale)
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(self.output_path, fourcc, fps, (out_w, out_h))
+
+        prev_t = time.time()
+        frame_count = 1
+        frame_up = cv2.resize(
+            first_frame_processed, (out_w, out_h), interpolation=cv2.INTER_LINEAR
+        )
+        prev_t = self._process_and_write_frame(
+            frame_up, out, prev_t, prediction_counter_df, frame_count
+        )
+
+        while True:
+            ok, frame_original_loop = cap.read()
+            if not ok:
+                break
+
+            if frame_count % (fps // 2 or 1) == 0:  # frame skipping @ ≈2 fps
+                frame_processed_loop = (
+                    cv2.rotate(frame_original_loop, cv2.ROTATE_90_CLOCKWISE)
+                    if frame_was_rotated
+                    else frame_original_loop
+                )
+                frame_up = cv2.resize(
+                    frame_processed_loop, (out_w, out_h), interpolation=cv2.INTER_LINEAR
+                )
+                prev_t = self._process_and_write_frame(
+                    frame_up, out, prev_t, prediction_counter_df, frame_count
+                )
+
+            frame_count += 1
+
+        cap.release()
+        out.release()
+        cv2.destroyAllWindows()
+        print(f"Processed {frame_count} frames. Finished → {self.output_path}")
+        return prediction_counter_df.fillna(0)
+
+    # ────────────────────────────────────────────────────────────────
+    def _process_and_write_frame(
+        self,
+        frame_up: np.ndarray,
+        out_writer: cv2.VideoWriter,
+        prev_t: float,
+        prediction_counter_df: pd.DataFrame,
+        frame_count: int,
+    ) -> float:
+        """Run YOLO on one frame, count per region, annotate, write, return timestamp."""
+        # Draw polygons first (scaled!)
+        scale_x = frame_up.shape[1] / (frame_up.shape[1] / self.scale)
+        scale_y = frame_up.shape[0] / (frame_up.shape[0] / self.scale)
+        for rid, poly in self.regions.items():
+            poly_up = (poly * [self.scale, self.scale]).astype(np.int32)
+            cv2.polylines(frame_up, [poly_up], True, (255, 255, 0), 2)
+            draw_text_with_background(frame_up, f"R{rid}", tuple(poly_up[0]), font_scale=0.8)
+
+        results = self.model.predict(
+            frame_up,
+            conf=self.conf,
+            classes=self.classes,
+            verbose=False,
+            device=self.device,
+        )
+
+        # counts[region][cls_id] → int
+        counts: Dict[int, Dict[int, int]] = {
+            rid: defaultdict(int) for rid in self.regions
+        }
+
+        if results and len(results[0].boxes):
+            xyxy = results[0].boxes.xyxy.cpu().numpy()
+            scores = results[0].boxes.conf.cpu().numpy()
+            cls_ids = results[0].boxes.cls.int().cpu().tolist()
+
+            for (x1, y1, x2, y2), score, cls_id in zip(xyxy, scores, cls_ids):
+                color = get_color_for_class(cls_id)
+                cv2.rectangle(
+                    frame_up, (int(x1), int(y1)), (int(x2), int(y2)), color, 2
+                )
+                label = LABEL_MAP.get(cls_id, f"Class {cls_id}")
+                draw_text_with_background(
+                    frame_up,
+                    f"{label}: {score:.2f}",
+                    (int(x1), int(y1) - 10),
+                    font_scale=0.6,
+                    font_thickness=1,
+                    bg_color=color,
+                    padding=3,
+                )
+
+                # Region assignment based on *centre* of the box
+                cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
+                for rid, poly in self.regions.items():
+                    poly_up = (poly * [self.scale, self.scale]).astype(np.int32)
+                    if _inside((cx, cy), poly_up):
+                        counts[rid][cls_id] += 1
+                        break  # one region per detection
+
+        # ─── Overlay per‑region counts + update DataFrame ───
+        df_idx = len(prediction_counter_df)
+        prediction_counter_df.at[df_idx, "Frame Number"] = frame_count
+
+        y_off = 30
+        for rid, cls_dict in counts.items():
+            for cls_id, cnt in cls_dict.items():
+                label = LABEL_MAP.get(cls_id, f"Class {cls_id}")
+                col_name = f"{label}_{rid}"
+                prediction_counter_df.at[df_idx, col_name] = cnt
+                draw_text_with_background(
+                    frame_up,
+                    f"{label}_{rid}: {cnt}",
+                    (10, y_off),
+                    font_scale=0.7,
+                    font_thickness=2,
+                    padding=6,
+                )
+                y_off += 25
+
+        # FPS overlay
+        now = time.time()
+        fps_live = 1.0 / (now - prev_t) if (now - prev_t) > 0 else 0.0
+        draw_text_with_background(
+            frame_up,
+            f"FPS: {fps_live:.1f}",
+            (10, frame_up.shape[0] - 20),
+            bg_color=(0, 0, 0),
+            text_color=(0, 255, 0),
+            font_scale=0.8,
+            padding=4,
+        )
+
+        out_writer.write(frame_up)
+        return now

streamlit_app.py

@@ -0,0 +1,280 @@
+# app.py
+
+import os
+import cv2
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from tempfile import NamedTemporaryFile
+from typing import List, Tuple, Dict
+
+import streamlit as st
+from PIL import Image
+from streamlit_drawable_canvas import st_canvas
+
+from lane_detection import YOLOVideoDetector, LABEL_MAP  # Your detector + LABEL_MAP from detection.py
+
+
+# ─────────────────────── Helper Functions ───────────────────────
+
+def extract_four_points(js) -> List[Tuple[int,int]]:
+    """
+    Return exactly four (x,y) clicks from streamlit_drawable_canvas JSON, or None.
+    """
+    if not js or "objects" not in js:
+        return None
+    pts = []
+    for obj in js["objects"]:
+        if obj.get("type") in {"circle", "rect"}:
+            x = int(obj["left"] + obj.get("radius", 0))
+            y = int(obj["top"] + obj.get("radius", 0))
+            pts.append((x, y))
+            if len(pts) == 4:
+                return pts
+    return None
+
+
+def draw_poly(img: np.ndarray, pts: List[Tuple[int,int]], color: Tuple[int,int,int]):
+    """
+    Draw a closed polygon (4 points) on img in the specified color.
+    """
+    cv2.polylines(img, [np.array(pts, np.int32)], True, color, 2, cv2.LINE_AA)
+
+
+# ───────────────────────── Streamlit App ─────────────────────────
+
+st.set_page_config(page_title="🚦 Multi‐Lane Congestion Demo", layout="wide")
+st.title("🚦 Multi‐Lane Vehicle Congestion Demo")
+
+# Initialize session state
+if "num_lanes" not in st.session_state:
+    st.session_state.num_lanes = None
+    st.session_state.current_lane = 0
+    st.session_state.lanes = []
+    st.session_state.video_path = None
+    st.session_state.video_uploaded = False
+
+# ─────────────────── Step 1: Choose Number of Lanes ───────────────────
+if st.session_state.num_lanes is None:
+    n = st.number_input(
+        "How many lanes would you like to define? (1–8)", 
+        min_value=1, 
+        max_value=8, 
+        value=2
+    )
+    if st.button("✔ Set Number of Lanes"):
+        st.session_state.num_lanes = int(n)
+        st.session_state.lanes = [None] * st.session_state.num_lanes
+    st.stop()
+
+# ─────────────────── Step 2: Upload a Video ───────────────────
+if not st.session_state.video_uploaded:
+    uploaded = st.file_uploader(
+        "Upload video (formats: mp4, avi, mov, mkv)", 
+        type=["mp4","avi","mov","mkv"]
+    )
+    if uploaded:
+        tmpfile = NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded.name)[1])
+        tmpfile.write(uploaded.read())
+        tmpfile.flush()
+        st.session_state.video_path = tmpfile.name
+        st.session_state.video_uploaded = True
+    else:
+        st.stop()
+
+# ─────────────────── Step 3: Grab First Frame & Scale ───────────────────
+cap = cv2.VideoCapture(st.session_state.video_path)
+ret, first_frame = cap.read()
+cap.release()
+if not ret:
+    st.error("❌ Could not read the first frame of the video.")
+    st.stop()
+
+frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
+h_orig, w_orig = frame_rgb.shape[:2]
+
+# If the frame is wider than 800 px, scale down
+MAX_W = 800
+if w_orig > MAX_W:
+    scale = MAX_W / w_orig
+    disp_w = MAX_W
+    disp_h = int(h_orig * scale)
+    frame_disp = cv2.resize(frame_rgb, (disp_w, disp_h), interpolation=cv2.INTER_AREA)
+else:
+    scale = 1.0
+    disp_w, disp_h = w_orig, h_orig
+    frame_disp = frame_rgb.copy()
+
+# ─────────────────── Step 4: Draw 4 Points Per Lane ───────────────────
+colors = [
+    (0, 255, 0), (255, 0, 0), (0, 0, 255), (255, 255, 0),
+    (255, 0, 255), (0, 255, 255), (128, 255, 0), (255, 128, 0)
+]
+
+if st.session_state.current_lane < st.session_state.num_lanes:
+    idx = st.session_state.current_lane
+    color = colors[idx % len(colors)]
+    st.subheader(f"2️⃣ Click exactly 4 points for Lane #{idx+1}")
+    st.caption("Draw 4 small circles on the image, then press **Confirm Lane**.")
+
+    canvas = st_canvas(
+        fill_color="rgba(0,0,0,0)",
+        stroke_width=2,
+        stroke_color=f"#{color[2]:02X}{color[1]:02X}{color[0]:02X}",
+        background_image=Image.fromarray(frame_disp),
+        drawing_mode="point",
+        key=f"lane_canvas_{idx}",
+        height=disp_h,
+        width=disp_w,
+        update_streamlit=True
+    )
+
+    pts_scaled = extract_four_points(canvas.json_data)
+    if pts_scaled:
+        preview = frame_disp.copy()
+        draw_poly(preview, pts_scaled, color)
+        st.image(preview, caption=f"Preview – Lane {idx+1}", use_column_width=True)
+
+    if st.button(f"Confirm Lane {idx+1}"):
+        if pts_scaled and len(pts_scaled) == 4:
+            # Convert scaled points back to original resolution
+            orig_pts = [(int(x/scale), int(y/scale)) for (x,y) in pts_scaled]
+            st.session_state.lanes[idx] = orig_pts
+            st.session_state.current_lane += 1
+        else:
+            st.warning("⚠ Please click exactly 4 points.")
+    st.stop()
+
+# ─────────────────── Step 5: Display All Lane Polygons ───────────────────
+st.subheader("✅ All lanes defined:")
+confirm_img = frame_rgb.copy()
+for i, poly in enumerate(st.session_state.lanes):
+    c = colors[i % len(colors)]
+    draw_poly(confirm_img, poly, c)
+    cv2.putText(
+        confirm_img, f"L{i+1}", (poly[0][0], poly[0][1] - 10),
+        cv2.FONT_HERSHEY_SIMPLEX, 1.0, c, 2, cv2.LINE_AA
+    )
+st.image(confirm_img, caption="All lane regions overlaid", use_column_width=True)
+
+# ─────────────────── Step 6: Input Thresholds & Run Congestion Analysis ───────────────────
+st.subheader("🔧 Congestion Thresholds")
+col1, col2 = st.columns(2)
+with col1:
+    low_thresh = st.number_input(
+        "Green if PCE <",
+        min_value=0.0,
+        max_value=20.0,
+        value=3.5,
+        step=0.1,
+        format="%.1f",
+        help="Values below this will be colored green"
+    )
+with col2:
+    high_thresh = st.number_input(
+        "Red if PCE >",
+        min_value=0.0,
+        max_value=20.0,
+        value=6.5,
+        step=0.1,
+        format="%.1f",
+        help="Values above this will be colored red"
+    )
+st.caption("Values between green/red thresholds will be yellow.")
+
+if st.button("🚀 Run Congestion Analysis"):
+    out_tmp = NamedTemporaryFile(delete=False, suffix=".mp4").name
+
+    regions: Dict[int, List[Tuple[int,int]]] = {
+        i: st.session_state.lanes[i] for i in range(st.session_state.num_lanes)
+    }
+
+    # Instantiate YOLOVideoDetector with 4 positional args
+    detector = YOLOVideoDetector(
+        "Weights/last.pt",  # <-- replace with your actual .pt path
+        st.session_state.video_path,
+        out_tmp,
+        regions
+    )
+    # Assign optional attributes
+    detector.classes = list(LABEL_MAP.keys())
+    detector.conf    = 0.35
+    detector.scale   = 1.5
+
+    with st.spinner("Processing video—this may take a while..."):
+        df = detector.process_video()
+
+    st.success("✅ Detection + annotation complete!")
+
+    # ───────────────── Compute Per-Lane PCE & Rolling Average ─────────────────
+    PCE = {
+        "auto":               0.8,
+        "bus":                4.0,
+        "car":                1.0,
+        "electric-rickshaw":  0.8,
+        "large-sized-truck":  4.5,
+        "medium-sized-truck": 3.5,
+        "motorbike":          0.5,
+        "small-sized-truck":  3.0,
+    }
+
+    for rid in regions.keys():
+        def lane_pce(row, rid=rid):
+            total = 0.0
+            for vt, factor in PCE.items():
+                coln = f"{vt}_{rid}"
+                cnt = row.get(coln, 0)
+                total += int(cnt) * factor
+            return total
+
+        df[f"PCE_lane{rid}"] = df.apply(lane_pce, axis=1)
+        df[f"PCE_lane{rid}_avg"] = df[f"PCE_lane{rid}"].rolling(window=5, min_periods=1).mean()
+
+    # ─────────────────── Lane-Wise Subplots with Smooth Line + Colored Markers ───────────────────
+    num_lanes = len(regions)
+    fig, axes = plt.subplots(num_lanes, 1, figsize=(10, 3 * num_lanes), sharex=True)
+
+    if num_lanes == 1:
+        axes = [axes]
+
+    for rid, ax in zip(regions.keys(), axes):
+        x = df["Frame Number"].values
+        y = df[f"PCE_lane{rid}_avg"].values
+
+        # 1) Plot a smooth continuous line in dark gray
+        ax.plot(x, y, color="gray", linewidth=1.2)
+
+        # 2) Overlay colored markers at each frame
+        colors_list = []
+        for yi in y:
+            if yi < low_thresh:
+                colors_list.append("green")
+            elif yi > high_thresh:
+                colors_list.append("red")
+            else:
+                colors_list.append("yellow")
+
+        ax.scatter(x, y, c=colors_list, s=20, edgecolors="black", linewidths=0.3)
+
+        ax.set_title(f"Lane {rid} PCE (rolling average)")
+        ax.set_ylabel("PCE")
+        ax.grid(alpha=0.3)
+
+    axes[-1].set_xlabel("Frame Number")
+    plt.tight_layout()
+
+    st.subheader("📊 Lane‐Wise Congestion Plots")
+    st.pyplot(fig)
+
+    # ─────────────────── Display Annotated Video & CSV Download ───────────────────
+    st.subheader("🎬 Annotated Output Video")
+    with open(out_tmp, "rb") as f:
+        st.video(f.read())
+
+    csv_bytes = df.to_csv(index=False).encode("utf-8")
+    st.download_button(
+        label="Download full counts + PCE CSV",
+        data=csv_bytes,
+        file_name="counts_and_pce.csv",
+        mime="text/csv"
+    )