Files
managed-portal/managed/store_dwell_alert/app/main.py
2026-04-27 10:04:36 +08:00

149 lines
5.1 KiB
Python

from __future__ import annotations
from argparse import ArgumentParser
from datetime import datetime
from pathlib import Path
from time import sleep
from app.config import (
AppConfig,
load_config,
resolve_config_path,
resolve_project_path,
resolve_project_root,
)
from app.modules.detector_tracker import YOLOTrackerAdapter
from app.modules.dwell_engine import DwellEngine
from app.modules.identity_resolver import IdentityResolver
from app.modules.notifier import append_json_event
from app.modules.staff_filter import StaffMatcher, load_staff_gallery
from app.modules.stream_reader import RTSPFrameReader
def build_app(config_path: str | Path | None = None) -> dict:
resolved_config_path = resolve_config_path(config_path)
config = load_config(resolved_config_path)
project_root = resolve_project_root(resolved_config_path)
event_sink_path = resolve_project_path(project_root, config.event_sink.path)
gallery_dir = resolve_project_path(project_root, config.staff.gallery_dir)
return {
"config": config,
"config_path": resolved_config_path,
"stream_reader": RTSPFrameReader(
rtsp_url=config.stream.rtsp_url,
sample_fps=config.stream.sample_fps,
reconnect_backoff_seconds=config.stream.reconnect_backoff_seconds,
),
"tracker": YOLOTrackerAdapter(),
"identity_resolver": IdentityResolver(
pause_timeout_seconds=config.thresholds.pause_timeout_seconds,
),
"staff_matcher": StaffMatcher(
gallery=load_staff_gallery(gallery_dir),
similarity_threshold=config.staff.similarity_threshold,
min_hits=config.staff.min_hits,
),
"dwell_engine": DwellEngine(
camera_id=config.camera_id,
min_people=config.thresholds.min_people,
min_dwell_seconds=config.thresholds.min_dwell_seconds,
pause_timeout_seconds=config.thresholds.pause_timeout_seconds,
alert_cooldown_seconds=config.thresholds.alert_cooldown_seconds,
),
"notifier": append_json_event,
"event_sink_path": event_sink_path,
}
def process_observations(
app: dict,
observations: list[dict],
when,
) -> list[dict]:
events = app["dwell_engine"].process_observations(observations, when)
for event in events:
app["notifier"](app["event_sink_path"], event)
return events
def process_frame(app: dict, frame, when: datetime | None = None) -> list[dict]:
observation_time = when or datetime.now().astimezone()
tracks = app["tracker"].track(frame)
observations = app["identity_resolver"].resolve(tracks, observation_time)
observations = app["staff_matcher"].classify(observations)
return process_observations(app, observations, observation_time)
def run_once(app: dict) -> list[dict]:
frame = app["stream_reader"].read()
if frame is None:
return []
return process_frame(app, frame)
def run_forever(app: dict, max_frames: int | None = None) -> int:
processed_frames = 0
try:
while True:
frame = app["stream_reader"].read()
if frame is None:
sleep(0.2)
continue
process_frame(app, frame)
processed_frames += 1
if max_frames is not None and processed_frames >= max_frames:
break
finally:
app["stream_reader"].close()
return processed_frames
def parse_args() -> ArgumentParser:
parser = ArgumentParser(description="Store dwell alert service bootstrap")
parser.add_argument("--config", help="Path to YAML config file")
parser.add_argument("--once", action="store_true", help="Read and process one frame")
parser.add_argument(
"--manage-api",
action="store_true",
help="Start the management API instead of the RTSP worker loop",
)
parser.add_argument("--host", default="0.0.0.0", help="Host for the management API")
parser.add_argument("--port", type=int, default=18081, help="Port for the management API")
parser.add_argument(
"--max-frames",
type=int,
help="Maximum frames to process before exiting",
)
return parser
def main() -> int:
parser = parse_args()
args = parser.parse_args()
if args.manage_api:
from app.manage_api import run_manage_api
run_manage_api(args.config, host=args.host, port=args.port)
return 0
app = build_app(args.config)
config: AppConfig = app["config"]
print(f"Loaded config from {app['config_path']}")
print(f"Camera: {config.camera_id}")
print(f"RTSP: {config.stream.rtsp_url}")
print(f"Event sink: {app['event_sink_path']}")
print(f"Loaded {len(app['staff_matcher'].gallery)} staff gallery embedding(s)")
if args.once:
events = run_once(app)
print(f"Generated {len(events)} event(s)")
app["stream_reader"].close()
return 0
processed_frames = run_forever(app, max_frames=args.max_frames)
print(f"Processed {processed_frames} frame(s)")
return 0
if __name__ == "__main__":
raise SystemExit(main())