File size: 2,948 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 | """Builds the Evidence Graph from raw detections — the single source of truth.
Associates each person to the vehicle their box overlaps most (rider↔motorcycle,
driver↔car); unassociated people become pedestrians.
"""
from __future__ import annotations
from core.schemas import (
DetectionResult,
Edge,
EvidenceGraph,
Light,
Person,
PersonRole,
Plate,
Vehicle,
)
VEHICLE_LABELS = {"car", "motorcycle", "motorbike", "truck", "bus", "bicycle"}
TWO_WHEELER = {"motorcycle", "bicycle"}
def _norm(label: str) -> str:
return "motorcycle" if label == "motorbike" else label
def build_graph(
image_id: str, det: DetectionResult, assoc_threshold: float = 0.15
) -> EvidenceGraph:
vehicles: list[Vehicle] = [
Vehicle(id=d.id, type=_norm(d.label), bbox=d.bbox, confidence=d.confidence)
for d in det.detections
if d.label in VEHICLE_LABELS
]
lights: list[Light] = [
Light(id=d.id, bbox=d.bbox, confidence=d.confidence)
for d in det.detections
if d.label == "traffic light"
]
persons: list[Person] = []
edges: list[Edge] = []
for d in det.detections:
if d.label != "person":
continue
best: Vehicle | None = None
best_score = 0.0
for v in vehicles:
score = d.bbox.intersection_over_self(v.bbox)
if score > best_score:
best_score, best = score, v
if best is not None and best_score >= assoc_threshold:
is_two = best.type in TWO_WHEELER
role = PersonRole.rider if is_two else PersonRole.driver
edge_type = "rides" if is_two else "drives"
persons.append(
Person(id=d.id, role=role, bbox=d.bbox, confidence=d.confidence)
)
edges.append(Edge(type=edge_type, src=d.id, dst=best.id))
else:
persons.append(
Person(
id=d.id,
role=PersonRole.pedestrian,
bbox=d.bbox,
confidence=d.confidence,
)
)
return EvidenceGraph(
image_id=image_id,
vehicles=vehicles,
persons=persons,
lights=lights,
edges=edges,
)
def attach_plates(
graph: EvidenceGraph, plates: list[Plate], threshold: float = 0.5
) -> None:
"""Attach each detected plate to the vehicle whose box it sits inside."""
for plate in plates:
graph.plates.append(plate)
if plate.bbox is None:
continue
best: Vehicle | None = None
best_score = 0.0
for v in graph.vehicles:
score = plate.bbox.intersection_over_self(v.bbox)
if score > best_score:
best_score, best = score, v
if best is not None and best_score >= threshold:
graph.edges.append(Edge(type="has_plate", src=best.id, dst=plate.id))
|