File size: 1,821 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 | """Scene-graph association: riders link to the motorcycle they overlap."""
from __future__ import annotations
from core.graph import build_graph
from core.schemas import BBox, Detection, DetectionResult, PersonRole
def _det(
label: str, box: tuple[float, float, float, float], conf: float = 0.9
) -> Detection:
return Detection(
label=label,
bbox=BBox(x1=box[0], y1=box[1], x2=box[2], y2=box[3]),
confidence=conf,
)
def test_rider_associates_to_overlapping_motorcycle() -> None:
det = DetectionResult(
image_width=200,
image_height=200,
detections=[
_det("motorcycle", (10, 60, 60, 140)),
_det("person", (15, 20, 55, 130)), # sits on the bike, overlaps heavily
],
)
g = build_graph("img1", det)
assert len(g.vehicles) == 1
riders = g.riders_of(g.vehicles[0].id)
assert len(riders) == 1
assert riders[0].role == PersonRole.rider
def test_far_person_is_pedestrian() -> None:
det = DetectionResult(
image_width=400,
image_height=200,
detections=[
_det("motorcycle", (10, 60, 60, 140)),
_det("person", (300, 20, 340, 120)), # nowhere near the bike
],
)
g = build_graph("img2", det)
assert g.riders_of(g.vehicles[0].id) == []
assert g.persons[0].role == PersonRole.pedestrian
def test_three_riders_link_to_one_bike() -> None:
det = DetectionResult(
image_width=200,
image_height=200,
detections=[
_det("motorcycle", (10, 60, 90, 150)),
_det("person", (15, 20, 45, 130)),
_det("person", (40, 20, 70, 130)),
_det("person", (60, 20, 90, 130)),
],
)
g = build_graph("img3", det)
assert len(g.riders_of(g.vehicles[0].id)) == 3
|