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:
1 parent
163280078f
commit
fc20a521ee
5 files changed
+370
-29
No files matched your search
@@ -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/<job_id>", methods=["DELETE"])
|
||||
def api_delete_job(job_id):
|
||||
"""Delete a completed/failed/cancelled job."""
|
||||
|
||||
+49
-28
@@ -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
|
||||
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)
|
||||
|
||||
# Auto-select truck detector if model lacks "truck" class
|
||||
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")
|
||||
|
||||
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,8 +203,7 @@ class JobQueue:
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
job.results.append(
|
||||
JobResult(
|
||||
job.results.append(JobResult(
|
||||
model_name=model_cfg.filename,
|
||||
output_path=result.output_path,
|
||||
loading_count=result.loading_count,
|
||||
@@ -224,17 +212,50 @@ class JobQueue:
|
||||
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)..."
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# Force release GPU memory between models
|
||||
import gc
|
||||
gc.collect()
|
||||
try:
|
||||
import torch
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
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:
|
||||
|
||||
+263
@@ -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,
|
||||
)
|
||||
+29
-1
@@ -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) {
|
||||
|
||||
@@ -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 @@
|
||||
<p>Upload a video and select detection models to run analysis.</p>
|
||||
</div>
|
||||
|
||||
<div class="active-jobs-banner" id="active-jobs-banner" style="display:none">
|
||||
<span class="badge badge-info">▶</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">
|
||||
<input type="hidden" name="video_path" id="video-path-input" value="">
|
||||
|
||||
|
||||
Reference in new issue
Block a user