Files
managed-portal/managed/people_flow_project/src/people_flow/models.py

137 lines
3.4 KiB
Python

from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
@dataclass
class YoloConfig:
model_path: str = "yolo11n.pt"
tracker: str = "botsort.yaml"
conf: float = 0.35
iou: float = 0.5
imgsz: int = 1280
device: str = "cuda:0"
@dataclass
class CountingConfig:
line: tuple[float, float, float, float] = (0.1, 0.55, 0.9, 0.55)
line_mode: str = "normalized"
crossing_tolerance: float = 12.0
def to_pixel_line(
self, width: int, height: int
) -> tuple[float, float, float, float]:
x1, y1, x2, y2 = self.line
if self.line_mode == "pixel":
return x1, y1, x2, y2
return x1 * width, y1 * height, x2 * width, y2 * height
@dataclass
class AttributeConfig:
enabled: bool = False
sample_every_n_frames: int = 12
max_samples_per_track: int = 5
min_person_box_width: int = 80
min_person_box_height: int = 160
person_crop_padding: float = 0.15
detector_backend: str = "retinaface"
enforce_detection: bool = False
@dataclass
class OutputConfig:
save_video: bool = True
save_json: bool = True
save_csv: bool = True
draw_boxes: bool = True
draw_labels: bool = True
@dataclass
class RtspConfig:
sample_interval_seconds: float = 1.0
window_seconds: int = 1800
reconnect_delay_seconds: float = 5.0
stream_open_timeout_seconds: float = 10.0
idle_sleep_seconds: float = 0.05
output_subdir: str = "rtsp_stream"
@dataclass
class QueueConfig:
enabled: bool = True
area: tuple[float, float, float, float] = (0.0, 0.0, 1.0, 1.0)
area_mode: str = "normalized"
queue_time_threshold_seconds: int = 300
crowded_count_threshold: int = 5
normal_count_threshold: int = 2
pause_timeout_seconds: int = 5
source_id: str = "people_flow_queue"
def to_pixel_area(
self, width: int, height: int
) -> tuple[float, float, float, float]:
x1, y1, x2, y2 = self.area
if self.area_mode == "pixel":
return x1, y1, x2, y2
return x1 * width, y1 * height, x2 * width, y2 * height
@dataclass
class WebhookConfig:
url: str = ""
timeout_seconds: float = 5.0
event_log_path: str = "outputs/rtsp_stream/webhook_events.jsonl"
@dataclass
class RuntimeConfig:
rtsp_url: str = "rtsp://user:password@camera-ip:554/h264/ch1/main/av_stream"
output_dir: str = "outputs"
@dataclass
class AppConfig:
yolo: YoloConfig = field(default_factory=YoloConfig)
counting: CountingConfig = field(default_factory=CountingConfig)
attributes: AttributeConfig = field(default_factory=AttributeConfig)
output: OutputConfig = field(default_factory=OutputConfig)
rtsp: RtspConfig = field(default_factory=RtspConfig)
queue: QueueConfig = field(default_factory=QueueConfig)
webhook: WebhookConfig = field(default_factory=WebhookConfig)
runtime: RuntimeConfig = field(default_factory=RuntimeConfig)
config_path: Path | None = None
@dataclass
class TrackObservation:
track_id: int
bbox: tuple[int, int, int, int]
confidence: float
center: tuple[float, float]
@dataclass
class CrossingEvent:
track_id: int
direction: str
@dataclass
class AttributeVote:
age: int
age_bucket: str
gender: str
@dataclass
class TrackAttributeSummary:
track_id: int
age: int
age_bucket: str
gender: str
samples_used: int