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
+38
-24
No files matched your search
+19
-18
@@ -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)
|
||||||
|
|||||||
+19
-6
@@ -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,11 +477,18 @@ 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, [])
|
||||||
|
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):
|
for future in as_completed(futures):
|
||||||
try:
|
try:
|
||||||
all_detections.extend(future.result())
|
all_detections.extend(future.result())
|
||||||
@@ -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 = []
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user