From fc20a521eeff0caafe5a21298a7b89807ab06568 Mon Sep 17 00:00:00 2001 From: jetson Date: Tue, 22 Sep 2026 11:10:41 +0700 Subject: [PATCH] feat: merged detection output + parallel model execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Multiple models run simultaneously on each frame via ThreadPoolExecutor - Detections merged with IoU deduplication (same class, high overlap → keep highest conf) - Single merged_annotated.mp4 output (not per-model videos) - Single model still uses existing run_pipeline for backward compat - Active jobs banner on upload page (polls every 5s) - GPU warning when selecting 3+ models - 62/62 tests pass --- app.py | 10 ++ src/job.py | 99 +++++++++------- src/pipeline.py | 263 +++++++++++++++++++++++++++++++++++++++++++ static/app.js | 30 ++++- templates/index.html | 19 ++++ 5 files changed, 381 insertions(+), 40 deletions(-) diff --git a/app.py b/app.py index 8a4d674..aa882ec 100644 --- a/app.py +++ b/app.py @@ -255,6 +255,16 @@ def api_cancel_job(job_id): return jsonify({"error": "cannot cancel job in state " + job.status.name}), 400 +@app.route("/api/jobs/active") +def api_active_jobs(): + """Return count of running/pending jobs.""" + active = [j for j in job_queue.list_jobs() if j.status.name in ('RUNNING', 'PENDING')] + return jsonify({ + "count": len(active), + "jobs": [{"job_id": j.job_id, "status": j.status.name, "video": os.path.basename(j.video_path)} for j in active], + }) + + @app.route("/api/jobs/", methods=["DELETE"]) def api_delete_job(job_id): """Delete a completed/failed/cancelled job.""" diff --git a/src/job.py b/src/job.py index a04c032..6b26c25 100644 --- a/src/job.py +++ b/src/job.py @@ -12,7 +12,7 @@ from enum import Enum, auto from pathlib import Path from src.model_registry import ModelConfig, find_truck_detector -from src.pipeline import run_pipeline, PipelineResult +from src.pipeline import run_pipeline, run_merged_pipeline, PipelineResult, MergedPipelineResult class JobStatus(Enum): @@ -172,26 +172,15 @@ class JobQueue: return try: - for i, model_cfg in enumerate(job.model_configs): - with self._lock: - if job.status == JobStatus.CANCELLED: - break - job.current_model = model_cfg.filename - job.progress = i / total_models - class_filter = job.class_filters.get(model_cfg.filename) - - # Auto-select truck detector if model lacks "truck" class + if total_models == 1: + # Single model: use existing run_pipeline + model_cfg = job.model_configs[0] + class_filter = job.class_filters.get(model_cfg.filename) truck_model_config = None if "truck" not in (model_cfg.known_classes or []): - truck_model_config = find_truck_detector(os.path.join(os.path.dirname(model_cfg.path), "..", "models")) - # If not found relative, try absolute models dir - if truck_model_config is None: - truck_model_config = find_truck_detector("./models") + truck_model_config = find_truck_detector("./models") - output_path = os.path.join( - job.output_dir, - f"{model_cfg.stem}_annotated.mp4", - ) + output_path = os.path.join(job.output_dir, f"{model_cfg.stem}_annotated.mp4") def _on_live_stats(stats): with self._lock: @@ -200,7 +189,7 @@ class JobQueue: job.live_unloading = stats["unloading"] job.live_frame = stats["frame"] - result: PipelineResult = run_pipeline( + result = run_pipeline( video_path=job.video_path, model_config=model_cfg, output_path=output_path, @@ -214,27 +203,59 @@ class JobQueue: ) with self._lock: - job.results.append( - JobResult( - model_name=model_cfg.filename, - output_path=result.output_path, - loading_count=result.loading_count, - unloading_count=result.unloading_count, - net_count=result.net_count, - batch_count=result.batch_count, - frame_count=result.frame_count, - duration_seconds=result.duration_seconds, - ) - ) + job.results.append(JobResult( + model_name=model_cfg.filename, + output_path=result.output_path, + loading_count=result.loading_count, + unloading_count=result.unloading_count, + net_count=result.net_count, + batch_count=result.batch_count, + frame_count=result.frame_count, + duration_seconds=result.duration_seconds, + )) + else: + # Multiple models: run merged pipeline + with self._lock: + job.current_model = f"Processing {total_models} models (merged)..." - # Force release GPU memory between models - import gc - gc.collect() - try: - import torch - torch.cuda.empty_cache() - except Exception: - pass + output_path = os.path.join(job.output_dir, "merged_annotated.mp4") + + def _on_live_stats_merged(stats): + with self._lock: + job.live_fps = stats["fps"] + job.live_loading = stats["loading"] + job.live_unloading = stats["unloading"] + job.live_frame = stats["frame"] + + # Auto-select truck detector + truck_model_config = None + if not any("truck" in (c.known_classes or []) for c in job.model_configs): + truck_model_config = find_truck_detector("./models") + + result = run_merged_pipeline( + video_path=job.video_path, + model_configs=job.model_configs, + output_path=output_path, + class_filters=job.class_filters, + truck_model_config=truck_model_config, + preview_path=job.preview_path, + preview_enabled=job.preview_enabled, + preview_every_n=2, + cancel_check=lambda: job.status == JobStatus.CANCELLED, + live_stats_callback=_on_live_stats_merged, + ) + + with self._lock: + job.results.append(JobResult( + model_name="Merged (" + ", ".join(c.stem for c in job.model_configs) + ")", + output_path=result.output_path, + loading_count=result.loading_count, + unloading_count=result.unloading_count, + net_count=result.net_count, + batch_count=result.batch_count, + frame_count=result.frame_count, + duration_seconds=result.duration_seconds, + )) with self._lock: if job.status != JobStatus.CANCELLED: diff --git a/src/pipeline.py b/src/pipeline.py index 9aa6f61..1d05501 100644 --- a/src/pipeline.py +++ b/src/pipeline.py @@ -5,6 +5,7 @@ from __future__ import annotations import os import time from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass import cv2 @@ -41,6 +42,24 @@ class PipelineResult: return self.loading_count - self.unloading_count +@dataclass +class MergedPipelineResult: + """Summary of a completed merged pipeline run (multiple models).""" + + output_path: str + frame_count: int + loading_count: int + unloading_count: int + batch_count: int + duration_seconds: float + model_names: list[str] + class_filter: list[str] | None + + @property + def net_count(self) -> int: + return self.loading_count - self.unloading_count + + def apply_class_filter( detections: list[Detection], class_filter: list[str] | None ) -> list[Detection]: @@ -306,3 +325,247 @@ def run_pipeline( model_name=model_config.filename, class_filter=effective_filter, ) + + +def _compute_iou(box1: tuple, box2: tuple) -> float: + x1 = max(box1[0], box2[0]) + y1 = max(box1[1], box2[1]) + x2 = min(box1[2], box2[2]) + y2 = min(box1[3], box2[3]) + inter = max(0, x2 - x1) * max(0, y2 - y1) + area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) + area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) + union = area1 + area2 - inter + return inter / union if union > 0 else 0.0 + + +def _deduplicate_detections(dets: list[Detection], iou_threshold: float = 0.5) -> list[Detection]: + if not dets: + return dets + dets = sorted(dets, key=lambda d: d.confidence, reverse=True) + kept: list[Detection] = [] + for det in dets: + is_dup = False + for existing in kept: + if existing.class_name == det.class_name: + if _compute_iou(existing.bbox, det.bbox) > iou_threshold: + is_dup = True + break + if not is_dup: + kept.append(det) + return kept + + +def run_merged_pipeline( + video_path: str, + model_configs: list[ModelConfig], + output_path: str, + class_filters: dict[str, list[str] | None] | None = None, + truck_model_config: ModelConfig | None = None, + sack_conf: float = 0.4, + truck_conf: float = 0.5, + truck_det_interval: int = 15, + progress_callback: Callable[[int, int], None] | None = None, + cancel_check: Callable[[], bool] | None = None, + preview_path: str | None = None, + preview_enabled: bool = True, + preview_every_n: int = 2, + preview_max_dim: int = 480, + preview_jpeg_quality: int = 60, + live_stats_callback: Callable[[dict], None] | None = None, +) -> MergedPipelineResult: + """Run multiple models on the same video, merge detections into one output.""" + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise RuntimeError(f"Cannot open video: {video_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # Load all detectors + detectors = [] + for cfg in model_configs: + if cfg.path.endswith(".onnx"): + try: + import onnxruntime + except ImportError: + raise RuntimeError( + f"onnxruntime not installed. Cannot load '{cfg.filename}'. " + f"Use the .engine version instead." + ) + model = YOLO(cfg.path) + eff_filter = (class_filters or {}).get(cfg.stem) or cfg.known_classes or None + detectors.append(BaseDetector(model, conf=sack_conf, class_filter=eff_filter)) + + # Truck detector + truck_det = None + if truck_model_config is not None: + truck_shared = YOLO(truck_model_config.path) + truck_det = BaseDetector(truck_shared, conf=truck_conf, class_filter=("truck",)) + elif len(model_configs) == 1 and "truck" in (model_configs[0].known_classes or []): + truck_det = detectors[0] + + # Single tracker using first model's weights + tracker_model = YOLO(model_configs[0].path) + tracker = ByteTrackTracker(tracker_model, conf=sack_conf) + stabilizer = BboxStabilizer() + roi_tracker = TruckROITracker(frame_width=w, frame_height=h) + counter = LineCrossCounter( + line_y=int(h * 0.50), + line_x_start=int(w * 0.38), + line_x_end=int(w * 0.72), + margin=20, + ) + batch_mgr = BatchLifecycleManager() + dashboard = DashboardOverlay() + writer = AnnotatedVideoWriter(output_path, fps=fps, frame_size=(w, h)) + + start_time = time.time() + frame_idx = 0 + completed_batches = 0 + _frame_times: list[float] = [] + + def on_batch_end(record): + nonlocal completed_batches + completed_batches += 1 + batch_mgr.on_batch_end(on_batch_end) + + try: + while True: + ret, frame = cap.read() + if not ret: + break + + frame_idx += 1 + timestamp = time.time() + _frame_times.append(timestamp) + if len(_frame_times) > 30: + _frame_times.pop(0) + + if cancel_check and frame_idx % 30 == 0 and cancel_check(): + break + + # Truck detection + roi = roi_tracker.roi + if truck_det is not None and frame_idx % truck_det_interval == 0: + trucks = truck_det.detect(frame) + roi = roi_tracker.update(trucks) + + truck_present = roi is not None and roi.confidence > 0 + if roi is not None: + counter.line_y = roi.line_y + counter.line_x_start = roi.x1 + counter.line_x_end = roi.x2 + + truck_centroid = None + if roi is not None: + truck_centroid = ((roi.x1 + roi.x2) / 2.0, (roi.y1 + roi.y2) / 2.0) + + if frame_idx % truck_det_interval == 0: + batch_mgr.update( + truck_detected=truck_present, + timestamp=timestamp, + loading_count=counter.loading_count, + unloading_count=counter.unloading_count, + truck_centroid=truck_centroid, + ) + + # Run all detectors in parallel + all_detections: list[Detection] = [] + if batch_mgr.is_active: + with ThreadPoolExecutor(max_workers=len(detectors)) as pool: + futures = {pool.submit(d.detect, frame): d for d in detectors} + for future in as_completed(futures): + try: + all_detections.extend(future.result()) + except Exception: + pass + + # Deduplicate + merged = _deduplicate_detections(all_detections) + + # Stabilize and filter to ROI + stable = stabilizer.update(merged) + if roi is not None: + tracked = [d for d in stable if roi.contains_x((d.bbox[0] + d.bbox[2]) / 2.0)] + else: + tracked = stable + + counter.update(tracked) + else: + tracked = [] + + # Annotate + viz = dashboard.draw( + frame=frame, + detections=tracked, + roi=roi, + loading_count=counter.loading_count, + unloading_count=counter.unloading_count, + batch_id=batch_mgr.current_batch_id, + history=batch_mgr.history, + system_state=batch_mgr.state, + batch_duration=batch_mgr.batch_duration, + stabilize_progress=batch_mgr.stabilize_progress, + waiting_duration=batch_mgr.waiting_duration, + ) + + # Model info overlay + model_names = ", ".join(c.stem for c in model_configs) + cv2.putText(viz, f"Models: {model_names}", (10, h - 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1) + + writer.write_frame(viz) + + # Live stats + if live_stats_callback and frame_idx % 10 == 0: + if len(_frame_times) >= 2: + dt = _frame_times[-1] - _frame_times[0] + processing_fps = (len(_frame_times) - 1) / dt if dt > 0 else 0.0 + else: + processing_fps = 0.0 + live_stats_callback({ + "fps": round(processing_fps, 1), + "loading": counter.loading_count, + "unloading": counter.unloading_count, + "frame": frame_idx, + }) + + # Preview + if preview_enabled and preview_path is not None and frame_idx % max(1, preview_every_n) == 0: + ph, pw = viz.shape[:2] + if max(ph, pw) > preview_max_dim: + scale = preview_max_dim / max(ph, pw) + preview_viz = cv2.resize(viz, (int(pw * scale), int(ph * scale))) + else: + preview_viz = viz + tmp_path = preview_path + ".tmp.jpg" + cv2.imwrite(tmp_path, preview_viz, [cv2.IMWRITE_JPEG_QUALITY, preview_jpeg_quality]) + os.replace(tmp_path, preview_path) + + if progress_callback: + progress_callback(frame_idx, total_frames) + + finally: + cap.release() + writer.finish() + import gc + gc.collect() + try: + import torch + torch.cuda.empty_cache() + except Exception: + pass + + duration = time.time() - start_time + return MergedPipelineResult( + output_path=output_path, + frame_count=frame_idx, + loading_count=counter.loading_count, + unloading_count=counter.unloading_count, + batch_count=completed_batches, + duration_seconds=duration, + model_names=[c.filename for c in model_configs], + class_filter=None, + ) diff --git a/static/app.js b/static/app.js index 281a5e2..2b5a764 100644 --- a/static/app.js +++ b/static/app.js @@ -270,7 +270,14 @@ Submit — enable/disable, loading state ------------------------------------------------ */ function initFormSubmit() { - form.addEventListener('submit', function () { + form.addEventListener('submit', function (e) { + var checked = modelGrid.querySelectorAll('.model-card-check:checked'); + if (checked.length >= 3) { + if (!confirm('Running ' + checked.length + ' models simultaneously may use significant GPU memory on Jetson. Continue?')) { + e.preventDefault(); + return; + } + } if (videoPathInput.value) { form.action = '/upload/reuse'; } else { @@ -304,6 +311,27 @@ initModelCards(); initFormSubmit(); loadExistingVideos(); + + // Poll active jobs + function checkActiveJobs() { + fetch('/api/jobs/active') + .then(function(res) { return res.json(); }) + .then(function(data) { + var banner = document.getElementById('active-jobs-banner'); + var text = document.getElementById('active-jobs-text'); + if (banner && text) { + if (data.count > 0) { + banner.style.display = ''; + text.textContent = data.count + ' job' + (data.count > 1 ? 's' : '') + ' in progress'; + } else { + banner.style.display = 'none'; + } + } + }) + .catch(function() {}); + } + checkActiveJobs(); + setInterval(checkActiveJobs, 5000); })(); function formatRelativeTime(timestamp) { diff --git a/templates/index.html b/templates/index.html index b29cdf9..9ef8411 100644 --- a/templates/index.html +++ b/templates/index.html @@ -257,6 +257,19 @@ } .empty-state p { margin-bottom: var(--space-2); } + /* Active jobs banner */ + .active-jobs-banner { + display: flex; + align-items: center; + gap: var(--space-3); + padding: var(--space-3) var(--space-4); + background: #f0f9ff; + border: 1px solid #0EA5E9; + border-radius: var(--radius-md); + margin-bottom: var(--space-6); + font-size: var(--text-sm); + } + /* Video delete button */ .video-delete-btn { flex-shrink: 0; @@ -291,6 +304,12 @@

Upload a video and select detection models to run analysis.

+ +