File size: 1,068 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 | """Traffic-light state via HSV colour analysis."""
from __future__ import annotations
import numpy as np
from PIL import Image
from core.detect import classify_lights
from core.schemas import BBox, EvidenceGraph, Light, LightState
def _light_image(path, color: tuple[int, int, int]) -> str:
arr = np.zeros((60, 60, 3), dtype="uint8")
arr[10:50, 10:50] = color
Image.fromarray(arr).save(path)
return str(path)
def test_red_light_classified(tmp_path) -> None:
p = _light_image(tmp_path / "r.png", (255, 0, 0))
g = EvidenceGraph(
image_id="i",
lights=[Light(id="tl", bbox=BBox(x1=10, y1=10, x2=50, y2=50), confidence=0.9)],
)
classify_lights(g, p)
assert g.lights[0].state == LightState.red
def test_green_light_classified(tmp_path) -> None:
p = _light_image(tmp_path / "g.png", (0, 255, 0))
g = EvidenceGraph(
image_id="i",
lights=[Light(id="tl", bbox=BBox(x1=10, y1=10, x2=50, y2=50), confidence=0.9)],
)
classify_lights(g, p)
assert g.lights[0].state == LightState.green
|