feat: merged detection output + parallel model execution

- 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
This commit is contained in:
jetson committed 2026-09-22 11:10:41 +07:00
1 parent 163280078f
commit fc20a521ee
5 files changed
+381 -40

No files matched your search

+10
View File
@@ -255,6 +255,16 @@ def api_cancel_job(job_id):
return jsonify({"error": "cannot cancel job in state " + job.status.name}), 400 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/<job_id>", methods=["DELETE"]) @app.route("/api/jobs/<job_id>", methods=["DELETE"])
def api_delete_job(job_id): def api_delete_job(job_id):
"""Delete a completed/failed/cancelled job.""" """Delete a completed/failed/cancelled job."""
+60 -39
View File
@@ -12,7 +12,7 @@ from enum import Enum, auto
from pathlib import Path from pathlib import Path
from src.model_registry import ModelConfig, find_truck_detector 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): class JobStatus(Enum):
@@ -172,26 +172,15 @@ class JobQueue:
return return
try: try:
for i, model_cfg in enumerate(job.model_configs): if total_models == 1:
with self._lock: # Single model: use existing run_pipeline
if job.status == JobStatus.CANCELLED: model_cfg = job.model_configs[0]
break class_filter = job.class_filters.get(model_cfg.filename)
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
truck_model_config = None truck_model_config = None
if "truck" not in (model_cfg.known_classes or []): 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")) truck_model_config = find_truck_detector("./models")
# If not found relative, try absolute models dir
if truck_model_config is None:
truck_model_config = find_truck_detector("./models")
output_path = os.path.join( output_path = os.path.join(job.output_dir, f"{model_cfg.stem}_annotated.mp4")
job.output_dir,
f"{model_cfg.stem}_annotated.mp4",
)
def _on_live_stats(stats): def _on_live_stats(stats):
with self._lock: with self._lock:
@@ -200,7 +189,7 @@ class JobQueue:
job.live_unloading = stats["unloading"] job.live_unloading = stats["unloading"]
job.live_frame = stats["frame"] job.live_frame = stats["frame"]
result: PipelineResult = run_pipeline( result = run_pipeline(
video_path=job.video_path, video_path=job.video_path,
model_config=model_cfg, model_config=model_cfg,
output_path=output_path, output_path=output_path,
@@ -214,27 +203,59 @@ class JobQueue:
) )
with self._lock: with self._lock:
job.results.append( job.results.append(JobResult(
JobResult( model_name=model_cfg.filename,
model_name=model_cfg.filename, output_path=result.output_path,
output_path=result.output_path, loading_count=result.loading_count,
loading_count=result.loading_count, unloading_count=result.unloading_count,
unloading_count=result.unloading_count, net_count=result.net_count,
net_count=result.net_count, batch_count=result.batch_count,
batch_count=result.batch_count, frame_count=result.frame_count,
frame_count=result.frame_count, duration_seconds=result.duration_seconds,
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 output_path = os.path.join(job.output_dir, "merged_annotated.mp4")
import gc
gc.collect() def _on_live_stats_merged(stats):
try: with self._lock:
import torch job.live_fps = stats["fps"]
torch.cuda.empty_cache() job.live_loading = stats["loading"]
except Exception: job.live_unloading = stats["unloading"]
pass 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: with self._lock:
if job.status != JobStatus.CANCELLED: if job.status != JobStatus.CANCELLED:
+263
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import os import os
import time import time
from collections.abc import Callable from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass from dataclasses import dataclass
import cv2 import cv2
@@ -41,6 +42,24 @@ class PipelineResult:
return self.loading_count - self.unloading_count 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( def apply_class_filter(
detections: list[Detection], class_filter: list[str] | None detections: list[Detection], class_filter: list[str] | None
) -> list[Detection]: ) -> list[Detection]:
@@ -306,3 +325,247 @@ def run_pipeline(
model_name=model_config.filename, model_name=model_config.filename,
class_filter=effective_filter, 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,
)
+29 -1
View File
@@ -270,7 +270,14 @@
Submit — enable/disable, loading state Submit — enable/disable, loading state
------------------------------------------------ */ ------------------------------------------------ */
function initFormSubmit() { 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) { if (videoPathInput.value) {
form.action = '/upload/reuse'; form.action = '/upload/reuse';
} else { } else {
@@ -304,6 +311,27 @@
initModelCards(); initModelCards();
initFormSubmit(); initFormSubmit();
loadExistingVideos(); 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) { function formatRelativeTime(timestamp) {
+19
View File
@@ -257,6 +257,19 @@
} }
.empty-state p { margin-bottom: var(--space-2); } .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 button */
.video-delete-btn { .video-delete-btn {
flex-shrink: 0; flex-shrink: 0;
@@ -291,6 +304,12 @@
<p>Upload a video and select detection models to run analysis.</p> <p>Upload a video and select detection models to run analysis.</p>
</div> </div>
<div class="active-jobs-banner" id="active-jobs-banner" style="display:none">
<span class="badge badge-info">&#9654;</span>
<span id="active-jobs-text"></span>
<a href="/jobs" class="btn btn-ghost btn-sm">View All</a>
</div>
<form action="/upload" method="post" enctype="multipart/form-data" id="upload-form"> <form action="/upload" method="post" enctype="multipart/form-data" id="upload-form">
<input type="hidden" name="video_path" id="video-path-input" value=""> <input type="hidden" name="video_path" id="video-path-input" value="">