File size: 3,415 Bytes
1c0c94d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
"""Rule Engine is pure -> test it with hand-built graphs, no models or I/O."""

from __future__ import annotations

from core.rules import helmet_rule, run_rules, seatbelt_rule, triple_riding_rule
from core.schemas import BBox, Edge, EvidenceGraph, Person, PersonRole, Tier, Vehicle


def _bike(vid: str = "v1") -> Vehicle:
    return Vehicle(
        id=vid, type="motorcycle", bbox=BBox(x1=0, y1=0, x2=50, y2=80), confidence=0.9
    )


def _rider(pid: str, helmet: bool | None) -> Person:
    return Person(
        id=pid,
        role=PersonRole.rider,
        bbox=BBox(x1=5, y1=0, x2=45, y2=40),
        confidence=0.8,
        helmet=helmet,
    )


def _graph(persons: list[Person], vid: str = "v1") -> EvidenceGraph:
    return EvidenceGraph(
        image_id="img_test",
        vehicles=[_bike(vid)],
        persons=persons,
        edges=[Edge(type="rides", src=p.id, dst=vid) for p in persons],
    )


def test_helmet_violation_when_rider_has_no_helmet() -> None:
    g = _graph([_rider("p1", helmet=False)])
    cands = helmet_rule(g)
    assert len(cands) == 1
    assert cands[0].type == "HELMET_NON_COMPLIANCE"
    assert cands[0].tier == Tier.A
    assert cands[0].rule_score == 1.0


def test_no_helmet_violation_when_rider_compliant() -> None:
    g = _graph([_rider("p1", helmet=True)])
    assert helmet_rule(g) == []


def test_unknown_helmet_is_not_flagged() -> None:
    # helmet None = classifier couldn't tell => not a violation (no noise row)
    g = _graph([_rider("p1", helmet=None)])
    assert helmet_rule(g) == []


def test_triple_riding_when_three_riders() -> None:
    g = _graph([_rider("p1", True), _rider("p2", True), _rider("p3", True)])
    cands = triple_riding_rule(g)
    assert len(cands) == 1
    assert cands[0].type == "TRIPLE_RIDING"
    assert "v1" in cands[0].subjects


def test_no_triple_riding_with_two_riders() -> None:
    g = _graph([_rider("p1", True), _rider("p2", True)])
    assert triple_riding_rule(g) == []


def test_run_rules_aggregates() -> None:
    g = _graph([_rider("p1", False), _rider("p2", True), _rider("p3", True)])
    types = {c.type for c in run_rules(g)}
    assert types == {"HELMET_NON_COMPLIANCE", "TRIPLE_RIDING"}


# --- seatbelt (Tier B) ----------------------------------------------------- #
def _car_with_driver(seatbelt: bool | None) -> EvidenceGraph:
    car = Vehicle(
        id="c1", type="car", bbox=BBox(x1=0, y1=0, x2=100, y2=100), confidence=0.9
    )
    drv = Person(
        id="d1",
        role=PersonRole.driver,
        bbox=BBox(x1=10, y1=10, x2=60, y2=90),
        confidence=0.8,
        seatbelt=seatbelt,
    )
    return EvidenceGraph(
        image_id="i",
        vehicles=[car],
        persons=[drv],
        edges=[Edge(type="drives", src="d1", dst="c1")],
    )


def test_seatbelt_candidate_is_speculative_tier_b() -> None:
    cands = seatbelt_rule(_car_with_driver(None))
    assert len(cands) == 1
    assert cands[0].type == "SEATBELT_NON_COMPLIANCE"
    assert cands[0].tier == Tier.B
    assert cands[0].speculative is True


def test_no_seatbelt_candidate_without_a_driver() -> None:
    car = Vehicle(
        id="c1", type="car", bbox=BBox(x1=0, y1=0, x2=100, y2=100), confidence=0.9
    )
    assert seatbelt_rule(EvidenceGraph(image_id="i", vehicles=[car])) == []


def test_no_seatbelt_candidate_when_visibly_belted() -> None:
    assert seatbelt_rule(_car_with_driver(True)) == []