Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| from collections import defaultdict | |
| from typing import Any, Literal | |
| from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError, model_validator | |
| from .models import ActionEvent, Detection, FrameSample | |
| def _label(value: str) -> str: | |
| return value.strip().lower() | |
| def _dedupe_labels(labels: list[str]) -> list[str]: | |
| seen: set[str] = set() | |
| result: list[str] = [] | |
| for label in labels: | |
| normalized = " ".join(label.strip().split()) | |
| key = normalized.lower() | |
| if normalized and key not in seen: | |
| seen.add(key) | |
| result.append(normalized) | |
| return result | |
| class PresentCondition(BaseModel): | |
| label: str | |
| min_count: int = Field(default=1, ge=1) | |
| class CountCondition(BaseModel): | |
| label: str | |
| minimum: int = Field(default=1, ge=1, validation_alias=AliasChoices("min", "minimum", "min_count")) | |
| class NearCondition(BaseModel): | |
| a: str | |
| b: str | |
| max_gap_percent: float = Field( | |
| default=16.0, | |
| ge=0.0, | |
| validation_alias=AliasChoices("max_gap_percent", "max_distance"), | |
| ) | |
| def migrate_max_distance(cls, data: Any) -> Any: | |
| if isinstance(data, dict) and "max_distance" in data and "max_gap_percent" not in data: | |
| value = data["max_distance"] | |
| if isinstance(value, int | float) and value <= 1: | |
| return {**data, "max_gap_percent": value * 100.0} | |
| return data | |
| class FarCondition(BaseModel): | |
| a: str | |
| b: str | |
| min_gap_percent: float = Field(default=25.0, ge=0.0) | |
| class CooldownCondition(BaseModel): | |
| key: str | None = None | |
| seconds: float | None = Field(default=None, ge=0.0) | |
| minutes: float | None = Field(default=None, ge=0.0) | |
| def duration_seconds(self) -> float: | |
| if self.seconds is not None: | |
| return self.seconds | |
| if self.minutes is not None: | |
| return self.minutes * 60.0 | |
| return 0.0 | |
| class ConditionBlock(BaseModel): | |
| present: PresentCondition | None = None | |
| count: CountCondition | None = None | |
| near: NearCondition | None = None | |
| far: FarCondition | None = None | |
| cooldown: CooldownCondition | None = None | |
| def exactly_one_condition(self) -> "ConditionBlock": | |
| selected = [self.present, self.count, self.near, self.far, self.cooldown] | |
| if sum(item is not None for item in selected) != 1: | |
| raise ValueError("Each condition block must contain exactly one condition.") | |
| return self | |
| class WhenClause(BaseModel): | |
| model_config = ConfigDict(populate_by_name=True) | |
| all_conditions: list[ConditionBlock] = Field(default_factory=list, alias="all") | |
| any_conditions: list[ConditionBlock] = Field(default_factory=list, alias="any") | |
| def has_conditions(self) -> "WhenClause": | |
| if not self.all_conditions and not self.any_conditions: | |
| raise ValueError("A rule needs at least one condition.") | |
| return self | |
| class ActionSpec(BaseModel): | |
| type: Literal["simulate", "webhook"] = "simulate" | |
| name: str | |
| url: str | None = None | |
| payload: dict[str, Any] = Field(default_factory=dict) | |
| class TriggerClause(BaseModel): | |
| on: Literal["while", "enter", "exit", "change"] = "enter" | |
| class ActionSet(BaseModel): | |
| model_config = ConfigDict(populate_by_name=True) | |
| enter: list[ActionSpec] = Field(default_factory=list) | |
| exit: list[ActionSpec] = Field(default_factory=list) | |
| while_actions: list[ActionSpec] = Field(default_factory=list, alias="while") | |
| class GateClause(BaseModel): | |
| enabled: bool = True | |
| cooldown: CooldownCondition | None = None | |
| class AutomationRule(BaseModel): | |
| name: str | |
| when: WhenClause | |
| trigger: TriggerClause = Field(default_factory=TriggerClause) | |
| gate: GateClause = Field(default_factory=GateClause) | |
| then: list[ActionSpec] | ActionSet | |
| def has_actions_for_trigger(self) -> "AutomationRule": | |
| if isinstance(self.then, list): | |
| if not self.then: | |
| raise ValueError("A rule needs at least one action.") | |
| if self.trigger.on == "change": | |
| raise ValueError("change triggers need then.enter/then.exit actions.") | |
| return self | |
| actions = _actions_for_trigger(self.then, self.trigger.on) | |
| if not actions: | |
| raise ValueError(f"A rule with trigger.on={self.trigger.on!r} needs matching actions.") | |
| return self | |
| class AutomationDocument(BaseModel): | |
| rules: list[AutomationRule] = Field(default_factory=list) | |
| def load_automation_text(text: str) -> AutomationDocument: | |
| """Load a JSON/YAML automation document and validate it.""" | |
| raw = text.strip() | |
| if not raw: | |
| raise ValueError("Rules are empty.") | |
| try: | |
| data = json.loads(raw) | |
| except json.JSONDecodeError: | |
| try: | |
| import yaml | |
| except ImportError as exc: # pragma: no cover - dependency guard | |
| raise RuntimeError("Install PyYAML to load YAML automation rules.") from exc | |
| data = yaml.safe_load(raw) | |
| if isinstance(data, list): | |
| data = {"rules": data} | |
| if isinstance(data, dict) and "rules" not in data and "name" in data: | |
| data = {"rules": [data]} | |
| try: | |
| return AutomationDocument.model_validate(data) | |
| except ValidationError: | |
| raise | |
| except Exception as exc: | |
| raise ValueError(f"Invalid automation document: {exc}") from exc | |
| def automation_schema() -> dict[str, Any]: | |
| return AutomationDocument.model_json_schema() | |
| def rule_labels(rule: AutomationRule) -> list[str]: | |
| labels: list[str] = [] | |
| for condition in [*rule.when.all_conditions, *rule.when.any_conditions]: | |
| if condition.present: | |
| labels.append(condition.present.label) | |
| if condition.count: | |
| labels.append(condition.count.label) | |
| if condition.near: | |
| labels.extend([condition.near.a, condition.near.b]) | |
| if condition.far: | |
| labels.extend([condition.far.a, condition.far.b]) | |
| return _dedupe_labels(labels) | |
| def document_labels(document: AutomationDocument, *, enabled_only: bool = True) -> list[str]: | |
| labels: list[str] = [] | |
| for rule in document.rules: | |
| if enabled_only and not rule.gate.enabled: | |
| continue | |
| labels.extend(rule_labels(rule)) | |
| return _dedupe_labels(labels) | |
| class RuleEngine: | |
| def __init__( | |
| self, | |
| rules: list[AutomationRule], | |
| last_fired: dict[str, float] | None = None, | |
| last_matched: dict[str, bool] | None = None, | |
| ) -> None: | |
| self.rules = rules | |
| self.last_fired: dict[str, float] = dict(last_fired or {}) | |
| self.last_matched: dict[str, bool] = dict(last_matched or {}) | |
| def evaluate_frame( | |
| self, | |
| detections: list[Detection], | |
| *, | |
| frame_index: int, | |
| timestamp_sec: float, | |
| ) -> list[ActionEvent]: | |
| events: list[ActionEvent] = [] | |
| for rule in self.rules: | |
| if not self._gate_allows(rule, timestamp_sec): | |
| self.last_matched[rule.name] = False | |
| continue | |
| matched = self._rule_matches(rule, detections, timestamp_sec) | |
| previous = self.last_matched.get(rule.name, False) | |
| edge = _trigger_edge(previous=previous, matched=matched) | |
| self.last_matched[rule.name] = matched | |
| actions = _actions_to_fire(rule, edge) | |
| if not matched and not actions: | |
| continue | |
| self._mark_cooldowns(rule, timestamp_sec) | |
| for action in actions: | |
| payload = { | |
| "rule": rule.name, | |
| "action": action.name, | |
| "trigger": edge, | |
| "frame_index": frame_index, | |
| "timestamp_sec": timestamp_sec, | |
| "detections": [item.model_dump(mode="json") for item in detections], | |
| } | |
| payload.update(action.payload) | |
| events.append( | |
| ActionEvent( | |
| rule=rule.name, | |
| action=action.name, | |
| type=action.type, | |
| frame_index=frame_index, | |
| timestamp_sec=timestamp_sec, | |
| url=action.url, | |
| payload=payload, | |
| ) | |
| ) | |
| return events | |
| def _rule_matches(self, rule: AutomationRule, detections: list[Detection], timestamp_sec: float) -> bool: | |
| all_ok = all( | |
| self._condition_matches(condition, detections, rule.name, timestamp_sec) | |
| for condition in rule.when.all_conditions | |
| ) | |
| any_ok = True | |
| if rule.when.any_conditions: | |
| any_ok = any( | |
| self._condition_matches(condition, detections, rule.name, timestamp_sec) | |
| for condition in rule.when.any_conditions | |
| ) | |
| return all_ok and any_ok | |
| def _gate_allows(self, rule: AutomationRule, timestamp_sec: float) -> bool: | |
| if not rule.gate.enabled: | |
| return False | |
| if not rule.gate.cooldown: | |
| return True | |
| return self._cooldown_allows(rule.gate.cooldown, rule.name, timestamp_sec) | |
| def _condition_matches( | |
| self, | |
| condition: ConditionBlock, | |
| detections: list[Detection], | |
| rule_name: str, | |
| timestamp_sec: float, | |
| ) -> bool: | |
| by_label = _group_by_label(detections) | |
| if condition.present: | |
| return len(by_label[_label(condition.present.label)]) >= condition.present.min_count | |
| if condition.count: | |
| return len(by_label[_label(condition.count.label)]) >= condition.count.minimum | |
| if condition.near: | |
| left = by_label[_label(condition.near.a)] | |
| right = by_label[_label(condition.near.b)] | |
| if not left or not right: | |
| return False | |
| return _min_box_gap_percent(left, right) <= condition.near.max_gap_percent | |
| if condition.far: | |
| left = by_label[_label(condition.far.a)] | |
| right = by_label[_label(condition.far.b)] | |
| if not left or not right: | |
| return False | |
| return _min_box_gap_percent(left, right) >= condition.far.min_gap_percent | |
| if condition.cooldown: | |
| return self._cooldown_allows(condition.cooldown, rule_name, timestamp_sec) | |
| return False | |
| def _mark_cooldowns(self, rule: AutomationRule, timestamp_sec: float) -> None: | |
| if rule.gate.cooldown: | |
| self.last_fired[rule.gate.cooldown.key or rule.name] = timestamp_sec | |
| for condition in [*rule.when.all_conditions, *rule.when.any_conditions]: | |
| if condition.cooldown: | |
| self.last_fired[condition.cooldown.key or rule.name] = timestamp_sec | |
| def _cooldown_allows(self, cooldown: CooldownCondition, rule_name: str, timestamp_sec: float) -> bool: | |
| key = cooldown.key or rule_name | |
| last = self.last_fired.get(key) | |
| return last is None or (timestamp_sec - last) >= cooldown.duration_seconds | |
| def evaluate_video_detections( | |
| rules: list[AutomationRule], | |
| detections: list[Detection], | |
| *, | |
| frames: list[FrameSample] | None = None, | |
| last_fired: dict[str, float] | None = None, | |
| ) -> tuple[list[ActionEvent], dict[str, float]]: | |
| engine = RuleEngine(rules, last_fired=last_fired) | |
| events: list[ActionEvent] = [] | |
| grouped: dict[tuple[int, float], list[Detection]] = defaultdict(list) | |
| for detection in detections: | |
| grouped[(detection.frame_index, detection.timestamp_sec)].append(detection) | |
| if frames: | |
| for frame in frames: | |
| grouped.setdefault((frame.frame_index, frame.timestamp_sec), []) | |
| for (frame_index, timestamp_sec), frame_detections in sorted(grouped.items()): | |
| events.extend( | |
| engine.evaluate_frame( | |
| frame_detections, | |
| frame_index=frame_index, | |
| timestamp_sec=timestamp_sec, | |
| ) | |
| ) | |
| return events, dict(engine.last_fired) | |
| def _group_by_label(detections: list[Detection]) -> dict[str, list[Detection]]: | |
| grouped: dict[str, list[Detection]] = defaultdict(list) | |
| for detection in detections: | |
| grouped[_label(detection.label)].append(detection) | |
| return grouped | |
| def _min_box_gap_percent(left: list[Detection], right: list[Detection]) -> float: | |
| best = float("inf") | |
| for left_detection in left: | |
| for right_detection in right: | |
| best = min(best, _box_gap_percent(left_detection, right_detection)) | |
| return best | |
| def _box_gap_percent(left: Detection, right: Detection) -> float: | |
| ax1, ay1, ax2, ay2 = left.bbox_xyxy_norm | |
| bx1, by1, bx2, by2 = right.bbox_xyxy_norm | |
| gap_x = max(0.0, max(bx1 - ax2, ax1 - bx2)) | |
| gap_y = max(0.0, max(by1 - ay2, ay1 - by2)) | |
| return max(gap_x, gap_y) * 100.0 | |
| def _trigger_edge(*, previous: bool, matched: bool) -> Literal["enter", "exit", "while", "none"]: | |
| if matched and not previous: | |
| return "enter" | |
| if not matched and previous: | |
| return "exit" | |
| if matched: | |
| return "while" | |
| return "none" | |
| def _actions_to_fire(rule: AutomationRule, edge: str) -> list[ActionSpec]: | |
| trigger = rule.trigger.on | |
| if trigger == "while": | |
| return _actions_for_trigger(rule.then, "while") if edge in {"enter", "while"} else [] | |
| if trigger == "enter": | |
| return _actions_for_trigger(rule.then, "enter") if edge == "enter" else [] | |
| if trigger == "exit": | |
| return _actions_for_trigger(rule.then, "exit") if edge == "exit" else [] | |
| if trigger == "change": | |
| return _actions_for_trigger(rule.then, edge) if edge in {"enter", "exit"} else [] | |
| return [] | |
| def _actions_for_trigger( | |
| actions: list[ActionSpec] | ActionSet, | |
| trigger: Literal["while", "enter", "exit", "change"], | |
| ) -> list[ActionSpec]: | |
| if isinstance(actions, list): | |
| return actions if trigger in {"while", "enter", "exit"} else [] | |
| if trigger == "while": | |
| return actions.while_actions or actions.enter | |
| if trigger == "enter": | |
| return actions.enter | |
| if trigger == "exit": | |
| return actions.exit | |
| return [*actions.enter, *actions.exit] | |