80 lines
3.1 KiB
Python
80 lines
3.1 KiB
Python
"""Per-frame detect + group + depth-gate + track + capture + emit pipeline.
|
|
|
|
Announcements are gated on two conditions:
|
|
* the track is currently stationary (``is_stationary()``), so walkers-by
|
|
never trigger a reject/TTS — only people planted at the checkpoint do;
|
|
* the current status differs from the last-announced status for the track,
|
|
so we never repeat the same announcement.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
from core.camera import estimate_person_distance_m
|
|
from core.capture import save_event_snapshot, save_track_image
|
|
from core.detection import collect_detections
|
|
from core.drawing import draw_track
|
|
from core.events import emit_event, write_result_csv
|
|
from core.grouping import group_detections_to_people
|
|
from core.paths import RESULT_CSV
|
|
from core.tracking import PersonTracker
|
|
|
|
|
|
def _filter_by_depth(candidates, depth_frame, depth_scale: float, max_distance_m: float):
|
|
"""Drop candidates whose median depth exceeds ``max_distance_m``.
|
|
|
|
Fail-open: candidates with no valid depth pass through. ``max_distance_m``
|
|
<= 0 disables the filter entirely.
|
|
"""
|
|
if depth_frame is None or max_distance_m <= 0 or not candidates:
|
|
return candidates
|
|
kept = []
|
|
for cand in candidates:
|
|
d = estimate_person_distance_m(depth_frame, cand.bbox, depth_scale)
|
|
if d is None or d <= max_distance_m:
|
|
cand.distance_m = d
|
|
kept.append(cand)
|
|
return kept
|
|
|
|
|
|
def process_frame(frame, model, tracker: PersonTracker, frame_idx: int, conf: float,
|
|
capture_dirs: Dict[str, Path], write_csv: bool = True,
|
|
*,
|
|
snapshot_dirs: Optional[Dict[str, Path]] = None,
|
|
depth_frame=None, depth_scale: float = 0.001,
|
|
max_distance_m: float = 0.0):
|
|
annotated = frame.copy()
|
|
h, w = annotated.shape[:2]
|
|
|
|
detections = collect_detections(frame, model, conf)
|
|
candidates = group_detections_to_people(detections, w, h)
|
|
candidates = _filter_by_depth(candidates, depth_frame, depth_scale, max_distance_m)
|
|
|
|
tracker.update(candidates, frame_idx)
|
|
visible = tracker.visible_tracks()
|
|
|
|
# Pass 1: save crops + draw annotations so snapshots capture the full scene.
|
|
for track in visible:
|
|
save_track_image(frame, track, capture_dirs)
|
|
draw_track(annotated, track)
|
|
|
|
# Pass 2: emit events ONLY for stationary tracks whose current status
|
|
# hasn't been announced yet. Walkers-by never reach this branch because
|
|
# their bbox-center history shows movement above the tolerance.
|
|
for track in visible:
|
|
if track.announced_status == track.status:
|
|
continue
|
|
if not track.is_stationary():
|
|
continue
|
|
ev_type = "NEW" if track.announced_status is None else "STATUS_CHANGE"
|
|
snap_path = None
|
|
if snapshot_dirs is not None:
|
|
snap_path = save_event_snapshot(annotated, track, snapshot_dirs)
|
|
emit_event(track, tracker.event_logger, ev_type, snapshot_path=snap_path)
|
|
|
|
if write_csv:
|
|
write_result_csv(list(tracker.tracks.values()), RESULT_CSV)
|
|
|
|
return annotated, visible
|