41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
"""Per-frame detect + group + track + capture + emit pipeline."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Dict
|
|
|
|
from saqr.core.capture import save_track_image
|
|
from saqr.core.detection import collect_detections
|
|
from saqr.core.drawing import draw_track
|
|
from saqr.core.events import emit_event, write_result_csv
|
|
from saqr.core.grouping import group_detections_to_people
|
|
from saqr.core.paths import RESULT_CSV
|
|
from saqr.core.tracking import PersonTracker
|
|
|
|
|
|
def process_frame(frame, model, tracker: PersonTracker, frame_idx: int, conf: float,
|
|
capture_dirs: Dict[str, Path], write_csv: bool = True):
|
|
annotated = frame.copy()
|
|
h, w = annotated.shape[:2]
|
|
|
|
detections = collect_detections(frame, model, conf)
|
|
candidates = group_detections_to_people(detections, w, h)
|
|
created, changed = tracker.update(candidates, frame_idx)
|
|
visible = tracker.visible_tracks()
|
|
|
|
created_ids = {t.track_id for t in created}
|
|
changed_ids = {t.track_id for t in changed}
|
|
event_ids = created_ids | changed_ids
|
|
|
|
for track in visible:
|
|
save_track_image(frame, track, capture_dirs)
|
|
if track.track_id in event_ids:
|
|
ev_type = "NEW" if track.track_id in created_ids else "STATUS_CHANGE"
|
|
emit_event(track, tracker.event_logger, ev_type)
|
|
draw_track(annotated, track)
|
|
|
|
if write_csv:
|
|
write_result_csv(list(tracker.tracks.values()), RESULT_CSV)
|
|
|
|
return annotated, visible
|