fix: use tracker for counting + throttle supplementary detectors

- Primary model uses ByteTrackTracker.update() every frame (gives track_ids)
- Supplementary models run every 5th frame only (reduces GPU load N->1)
- Counter only receives detections with track_id (fixes zero counts)
- Dashboard draws all detections (tracker + supplementary merged)
This commit is contained in:
jetson committed 2026-09-22 11:42:10 +07:00
1 parent fc20a521ee
commit 5bf44d6457
2 files changed
+43 -29

No files matched your search

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