diff --git a/src/dashboard.py b/src/dashboard.py index 152503a..0077d74 100644 --- a/src/dashboard.py +++ b/src/dashboard.py @@ -55,17 +55,16 @@ class DashboardOverlay: stabilize_progress: float = 0.0, waiting_duration: float = 0.0, ) -> np.ndarray: - out = frame.copy() if roi is not None: - self._draw_roi(out, roi) - self._draw_counting_zone(out, roi) - self._draw_detections(out, detections) - self._draw_stats(out, loading_count, unloading_count, batch_id, history) + self._draw_roi(frame, roi) + self._draw_counting_zone(frame, roi) + self._draw_detections(frame, detections) + self._draw_stats(frame, loading_count, unloading_count, batch_id, history) self._draw_system_state( - out, system_state, batch_duration, idle_timer, stabilize_progress, + frame, system_state, batch_duration, idle_timer, stabilize_progress, waiting_duration, ) - return out + return frame def _draw_roi(self, frame: np.ndarray, roi: TruckROI) -> None: cv2.rectangle( @@ -84,13 +83,15 @@ class DashboardOverlay: self, frame: np.ndarray, roi: TruckROI, margin: int = 20, ) -> None: y = roi.line_y - # Draw zone band (semi-transparent) - overlay = frame.copy() - cv2.rectangle( - overlay, (roi.x1, y - margin), (roi.x2, y + margin), - MAGENTA, -1, - ) - cv2.addWeighted(overlay, 0.15, frame, 0.85, 0, frame) + h, w = frame.shape[:2] + # Blend only the narrow strip region, not the full frame + sy = max(0, y - margin) + ey = min(h, y + margin) + sx = max(0, roi.x1) + ex = min(w, roi.x2) + strip = frame[sy:ey, sx:ex].copy() + cv2.rectangle(strip, (0, 0), (ex - sx, ey - sy), MAGENTA, -1) + cv2.addWeighted(strip, 0.15, frame[sy:ey, sx:ex], 0.85, 0, frame[sy:ey, sx:ex]) # Draw center line cv2.line(frame, (roi.x1, y), (roi.x2, y), MAGENTA, 2) cv2.putText( @@ -188,12 +189,12 @@ class DashboardOverlay: # Get display info for current state label, color = STATE_DISPLAY.get(system_state, ("UNKNOWN", GRAY)) - # Draw state bar background + # Draw state bar — blend only the bottom bar strip, not full frame bar_h = 36 bar_y = h - bar_h - overlay = frame.copy() - cv2.rectangle(overlay, (0, bar_y), (w, h), (0, 0, 0), -1) - cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame) + strip = frame[bar_y:h, 0:w].copy() + cv2.rectangle(strip, (0, 0), (w, bar_h), (0, 0, 0), -1) + cv2.addWeighted(strip, 0.7, frame[bar_y:h, 0:w], 0.3, 0, frame[bar_y:h, 0:w]) # Draw colored indicator dot cv2.circle(frame, (20, bar_y + bar_h // 2), 8, color, -1) diff --git a/src/pipeline.py b/src/pipeline.py index 1d05501..ccd3b54 100644 --- a/src/pipeline.py +++ b/src/pipeline.py @@ -404,8 +404,12 @@ def run_merged_pipeline( 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] + # ponytail: first-match truck detector; could prefer dedicated truck model over sack model w/ truck class + else: + for i, cfg in enumerate(model_configs): + if "truck" in (cfg.known_classes or []): + truck_det = detectors[i] + break # Single tracker using first model's weights tracker_model = YOLO(model_configs[0].path) @@ -426,6 +430,7 @@ def run_merged_pipeline( frame_idx = 0 completed_batches = 0 _frame_times: list[float] = [] + sup_det_interval = 5 # run supplementary models every N frames def on_batch_end(record): nonlocal completed_batches @@ -472,16 +477,23 @@ def run_merged_pipeline( truck_centroid=truck_centroid, ) - # Run all detectors in parallel + # Detection + tracking 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 + # Primary: tracker gives detections WITH track_id for counting + primary = tracker.update(frame, []) + all_detections.extend(primary) + + # Supplementary: run extra models every N frames for richer annotations + if len(detectors) > 1 and frame_idx % sup_det_interval == 0: + with ThreadPoolExecutor(max_workers=len(detectors) - 1) as pool: + extras = [d for i, d in enumerate(detectors) if i != 0] + futures = {pool.submit(d.detect, frame): d for d in extras} + for future in as_completed(futures): + try: + all_detections.extend(future.result()) + except Exception: + pass # Deduplicate merged = _deduplicate_detections(all_detections) @@ -493,7 +505,8 @@ def run_merged_pipeline( else: tracked = stable - counter.update(tracked) + # Count using only detections with track_id (tracker results) + counter.update([d for d in tracked if d.track_id is not None]) else: tracked = []