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:
1 parent
fc20a521ee
commit
5bf44d6457
2 files changed
+43
-29
No files matched your search
+19
-18
@@ -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)
|
||||
|
||||
+24
-11
@@ -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 = []
|
||||
|
||||
|
||||
Reference in new issue
Block a user