Merge branch 'main' of https://huggingface.co/spaces/Perception365/UrbanFlow

Dockerfile

@@ -9,7 +9,7 @@ RUN apt-get update && apt-get install -y \
     ffmpeg \
     && rm -rf /var/lib/apt/lists/*
 
-COPY requirements.txt .
+COPY requirements.txt . 
 RUN pip install --no-cache-dir -r requirements.txt
 
 COPY backend/ ./backend/

backend/engine.py

@@ -4,6 +4,7 @@ import tempfile
 import numpy as np
 import cv2
 from collections import defaultdict
+from constants import MODEL_CLASSES
 
 
 def _side(p, a, b):
@@ -25,27 +26,50 @@ _CLR_LINE = (80, 220, 100)    # green
 _CLR_TEXT_BG = (30, 30, 30)   # dark bg for text
 
 
-def _draw_annotations(frame, boxes, ids, line_pts):
+def _draw_annotations(frame, boxes, ids, clses, line_pts, options):
     """Draw bounding boxes, track IDs, and counting line on frame in-place."""
-    # Counting line
-    cv2.line(frame, tuple(line_pts[0]), tuple(line_pts[1]), _CLR_LINE, 3, cv2.LINE_AA)
+    # Counting line (Spatial Boundary)
+    if options.get("spatial", True):
+        cv2.line(frame, tuple(line_pts[0]), tuple(line_pts[1]), _CLR_LINE, 3, cv2.LINE_AA)
 
-    if boxes is not None and ids is not None:
-        for box, obj_id in zip(boxes, ids):
+    if boxes is not None and ids is not None and clses is not None:
+        for box, obj_id, cls_idx in zip(boxes, ids, clses):
             x1, y1, x2, y2 = map(int, box)
-            cv2.rectangle(frame, (x1, y1), (x2, y2), _CLR_BOX, 2)
-            label = f"ID:{int(obj_id)}"
-            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
-            cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw + 6, y1), _CLR_TEXT_BG, -1)
-            cv2.putText(frame, label, (x1 + 3, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA)
-
-
-def run(model, video_path, line, config, on_frame, save_annotated=False):
+            
+            # Bounding Box
+            if options.get("bbox", True):
+                cv2.rectangle(frame, (x1, y1), (x2, y2), _CLR_BOX, 2)
+
+            # Labels
+            labels = []
+            if options.get("track_id", True):
+                labels.append(f"ID:{int(obj_id)}")
+            if options.get("class_name", True):
+                labels.append(MODEL_CLASSES.get(int(cls_idx), "Unknown"))
+            elif options.get("class_id", False):
+                labels.append(f"C:{int(cls_idx)}")
+            
+            if labels:
+                label_text = " | ".join(labels)
+                (tw, th), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
+                cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw + 6, y1), _CLR_TEXT_BG, -1)
+                cv2.putText(frame, label_text, (x1 + 3, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA)
+
+
+def run(model, video_path, line, config, on_frame, save_annotated=False, annotated_options=None):
     """
     Runs YOLO tracking on video. Calls on_frame(update_dict) after each processed frame.
     line: [[x1,y1], [x2,y2]]
     save_annotated: if True, writes annotated MP4 with boxes + IDs + counting line
+    annotated_options: dict of toggleable visual overlays
     """
+    if annotated_options is None:
+        annotated_options = {"bbox": True, "track_id": True, "spatial": True}
+    
+    # Force bbox to True if export is enabled (user requirement)
+    if save_annotated:
+        annotated_options["bbox"] = True
+
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
     total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -53,6 +77,9 @@ def run(model, video_path, line, config, on_frame, save_annotated=False):
     out_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     cap.release()
 
+    # Dynamic crossing threshold: 5% of frame height, min 40px
+    cross_dist = max(40, int(out_h * 0.05))
+
     stride = config["detect_stride"]
     total_iters = total // stride
 
@@ -119,7 +146,7 @@ def run(model, video_path, line, config, on_frame, save_annotated=False):
                 if obj_id in prev_side and obj_id not in counted_ids:
                     if prev_side[obj_id] != current:
                         dist = _point_to_segment_dist(cx, cy, a[0], a[1], b[0], b[1])
-                        if dist < 12:
+                        if dist < cross_dist:
                             t = frame_idx * stride / fps
                             flow_times.append(round(t, 2))
 
@@ -135,7 +162,7 @@ def run(model, video_path, line, config, on_frame, save_annotated=False):
         # Write annotated frame (only for frames we already process)
         if writer is not None:
             frame = r.orig_img.copy()
-            _draw_annotations(frame, cur_boxes, cur_ids, [a, b])
+            _draw_annotations(frame, cur_boxes, cur_ids, cls if r.boxes.id is not None else None, [a, b], annotated_options)
             writer.write(frame)
 
         congestion.append(active)

backend/server.py

@@ -125,6 +125,7 @@ async def ws_run(ws: WebSocket):
     line = data["line"]
     cfg = data["config"]
     save_annotated = data.get("annotated_video", False)
+    annotated_options = data.get("annotated_options", {"bbox": True, "track_id": True, "spatial": True})
     report_format = data.get("report_format", "png")
 
     path = videos.get(video_id)
@@ -137,7 +138,7 @@ async def ws_run(ws: WebSocket):
         loop.call_soon_threadsafe(queue.put_nowait, update)
 
     task = loop.run_in_executor(
-        None, run, model, path, line, cfg, on_frame, save_annotated
+        None, run, model, path, line, cfg, on_frame, save_annotated, annotated_options
     )
 
     try:

frontend/initial.html

@@ -1,7 +1,20 @@
 <!DOCTYPE html>
-<html lang="en"> 
+<html lang="en">
 
 <head>
+    <!-- SPA Bootstrap: load dashboard without changing URL -->
+    <script>
+        if (sessionStorage.getItem('funky_run')) {
+            document.documentElement.style.display = 'none';
+            fetch('/vehicles.html')
+                .then(function (r) { return r.text(); })
+                .then(function (html) {
+                    document.open();
+                    document.write(html);
+                    document.close();
+                });
+        }
+    </script>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>UrbanFlow</title>
@@ -347,7 +360,7 @@
                 line: line,
                 config: runConfig
             }));
-            window.location.href = 'vehicles.html';
+            window.location.href = '/';
         }
     </script>
 </body>

frontend/vehicles.html

@@ -262,6 +262,68 @@
         #btn-start-processing {
             font-family: 'Montserrat', sans-serif !important;
         }
+
+        /* Chips */
+        .chip-container {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 8px;
+            margin-top: 12px;
+            padding-top: 12px;
+            border-top: 1px solid #1a1a1a;
+            transition: all 0.3s ease;
+        }
+
+        .chip {
+            display: inline-flex;
+            align-items: center;
+            gap: 6px;
+            padding: 6px 14px;
+            border-radius: 9999px;
+            font-size: 10px;
+            font-weight: 700;
+            cursor: pointer;
+            transition: all 0.2s ease;
+            user-select: none;
+            border: 1px solid #333333;
+            background: rgba(255, 255, 255, 0.03);
+            color: #888888;
+        }
+
+        .chip.active {
+            background: #ffffff;
+            color: #000000;
+            border-color: #ffffff;
+        }
+
+        .chip.frozen {
+            background: rgba(255, 255, 255, 0.4);
+            color: #000000;
+            border-color: transparent;
+            cursor: default !important;
+            pointer-events: none;
+        }
+
+        .chip:hover {
+            border-color: #666666;
+        }
+
+        .chip.active:hover {
+            background: #eeeeee;
+        }
+
+        .chip i {
+            font-size: 9px;
+        }
+
+        .hidden-chip-container {
+            height: 0;
+            opacity: 0;
+            overflow: hidden;
+            margin-top: 0;
+            padding-top: 0;
+            border-top: none;
+        }
     </style>
 </head>
 
@@ -611,14 +673,34 @@
                                 <option value="pdf">PDF Document</option>
                             </select>
                         </div>
-                        <div class="s-row" data-param="annotated">
-                            <div>
-                                <div class="text-xs font-semibold text-slate-700">Export Annotated Video</div>
-                                <div class="text-[10px] text-slate-400">Bounding boxes + track IDs + spatial boundary
-                                    overlays</div>
+                        <div class="s-row flex-col items-stretch" data-param="annotated">
+                            <div class="flex items-center justify-between w-full">
+                                <div>
+                                    <div class="text-xs font-semibold text-slate-700">Export Annotated Video</div>
+                                    <div class="text-[10px] text-slate-400">Layered visual overlays for diagnostic
+                                        analysis</div>
+                                </div>
+                                <div class="toggle-track" id="sv-annotated" onclick="toggleExportMaster(this)">
+                                    <div class="toggle-thumb"></div>
+                                </div>
                             </div>
-                            <div class="toggle-track" id="sv-annotated" onclick="this.classList.toggle('active')">
-                                <div class="toggle-thumb"></div>
+
+                            <div id="chip-selector" class="chip-container hidden-chip-container">
+                                <div class="chip frozen" id="chip-bbox">
+                                    <i class="fa-solid fa-check"></i> Bounding Boxes
+                                </div>
+                                <div class="chip active" id="chip-spatial" onclick="toggleChip('spatial')">
+                                    <i class="fa-solid fa-check"></i> Spatial Boundary
+                                </div>
+                                <div class="chip active" id="chip-class_id" onclick="toggleChip('class_id')">
+                                    <i class="fa-solid fa-check"></i> Class ID
+                                </div>
+                                <div class="chip active" id="chip-class_name" onclick="toggleChip('class_name')">
+                                    <i class="fa-solid fa-check"></i> Class Names
+                                </div>
+                                <div class="chip" id="chip-track_id" onclick="toggleChip('track_id')">
+                                    <i class="fa-solid fa-plus"></i> Track IDs
+                                </div>
                             </div>
                         </div>
                     </div>
@@ -1121,6 +1203,15 @@
             const reportFmt = document.getElementById('sv-report').value;
             const annotated = document.getElementById('sv-annotated').classList.contains('active');
 
+            // Annotation Options
+            const annotated_options = {
+                bbox: true, // Always true if export is enabled
+                spatial: document.getElementById('chip-spatial').classList.contains('active'),
+                class_name: document.getElementById('chip-class_name').classList.contains('active'),
+                class_id: document.getElementById('chip-class_id').classList.contains('active'),
+                track_id: document.getElementById('chip-track_id').classList.contains('active')
+            };
+
             // Apply to config
             _params.config.imgsz = imgsz;
             _params.config.conf = conf;
@@ -1158,6 +1249,7 @@
                     line: _params.line,
                     config: _params.config,
                     annotated_video: annotated,
+                    annotated_options: annotated_options,
                     report_format: reportFmt
                 }));
             };
@@ -1236,7 +1328,7 @@
             'class_dominance.pdf': { title: 'Class Dominance', desc: 'Vehicle count by classification type' },
             'confidence_dist.png': { title: 'Confidence Distribution', desc: 'Detection confidence histogram' },
             'confidence_dist.pdf': { title: 'Confidence Distribution', desc: 'Detection confidence histogram' },
-            'annotated.mp4': { title: 'Annotated Video', desc: 'Video with bounding boxes, track IDs, and spatial boundary overlays' },
+            'annotated.mp4': { title: 'Annotated Video', desc: 'Video with custom diagnostic overlays (Boxes, Names, IDs, etc.)' },
         };
 
         async function loadReports(videoId) {
@@ -1293,6 +1385,27 @@
             });
         }
 
+        function toggleExportMaster(el) {
+            el.classList.toggle('active');
+            const chips = document.getElementById('chip-selector');
+            if (el.classList.contains('active')) {
+                chips.classList.remove('hidden-chip-container');
+            } else {
+                chips.classList.add('hidden-chip-container');
+            }
+        }
+
+        function toggleChip(id) {
+            const chip = document.getElementById(`chip-${id}`);
+            chip.classList.toggle('active');
+            const icon = chip.querySelector('i');
+            if (chip.classList.contains('active')) {
+                icon.className = 'fa-solid fa-check';
+            } else {
+                icon.className = 'fa-solid fa-plus';
+            }
+        }
+
         init();
     </script>
 </body>