70 lines
2.8 KiB
Python
70 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from yovision_brain.decode import DecodedFrame
|
|
from yovision_brain.vision import (
|
|
BoundingBox,
|
|
Detection,
|
|
LumaBlobDetector,
|
|
SingleStreamTracker,
|
|
TorchLumaBlobDetector,
|
|
)
|
|
|
|
|
|
def frame(payload: bytes, *, sequence: int = 0, width: int = 4, height: int = 3) -> DecodedFrame:
|
|
return DecodedFrame(sequence, sequence * 40_000_000, "camera", "main", width, height, "rgb24", payload)
|
|
|
|
|
|
def rgb(values: list[int]) -> bytes:
|
|
return b"".join(bytes((value, value, value)) for value in values)
|
|
|
|
|
|
def detection(left: int, top: int, right: int, bottom: int) -> Detection:
|
|
return Detection(BoundingBox(left, top, right, bottom), "anonymous_target", 0.9)
|
|
|
|
|
|
def test_detector_emits_only_anonymous_observations() -> None:
|
|
payload = rgb([0, 255, 255, 0, 0, 255, 255, 0, 0, 0, 0, 0])
|
|
result = LumaBlobDetector(minimum_area=2).detect(frame(payload))
|
|
assert result == (Detection(BoundingBox(1, 0, 3, 2), "anonymous_target", 1.0),)
|
|
assert LumaBlobDetector.metadata.weights == "none"
|
|
assert "external model license" in LumaBlobDetector.metadata.license
|
|
|
|
|
|
def test_empty_frame_has_no_detection() -> None:
|
|
assert LumaBlobDetector().detect(frame(rgb([0] * 12))) == ()
|
|
|
|
|
|
def test_tracker_keeps_session_id_across_motion_and_short_occlusion() -> None:
|
|
tracker = SingleStreamTracker(iou_threshold=0.1, max_missed=2)
|
|
first = tracker.update((detection(0, 0, 3, 3),), frame_sequence=0, timestamp_ns=0)
|
|
assert first[0].track_id == "track-000001"
|
|
assert tracker.update((), frame_sequence=1, timestamp_ns=1) == ()
|
|
resumed = tracker.update((detection(1, 0, 4, 3),), frame_sequence=2, timestamp_ns=2)
|
|
assert resumed[0].track_id == "track-000001"
|
|
assert tracker.finish() == ("track-000001",)
|
|
|
|
|
|
def test_disappeared_track_ends_and_new_target_gets_new_id() -> None:
|
|
tracker = SingleStreamTracker(max_missed=1)
|
|
first = tracker.update((detection(0, 0, 2, 2),), frame_sequence=0, timestamp_ns=0)
|
|
tracker.update((), frame_sequence=1, timestamp_ns=1)
|
|
tracker.update((), frame_sequence=2, timestamp_ns=2)
|
|
second = tracker.update((detection(0, 0, 2, 2),), frame_sequence=3, timestamp_ns=3)
|
|
assert first[0].track_id == "track-000001"
|
|
assert second[0].track_id == "track-000002"
|
|
|
|
|
|
def test_track_ids_are_session_local() -> None:
|
|
one = SingleStreamTracker().update((detection(0, 0, 1, 1),), frame_sequence=0, timestamp_ns=0)
|
|
two = SingleStreamTracker().update((detection(0, 0, 1, 1),), frame_sequence=0, timestamp_ns=0)
|
|
assert one[0].track_id == two[0].track_id == "track-000001"
|
|
|
|
|
|
def test_torch_backend_cpu_smoke_uses_no_external_weights() -> None:
|
|
pytest.importorskip("torch")
|
|
result = TorchLumaBlobDetector().detect(frame(rgb([0, 255] + [0] * 10)))
|
|
assert result[0].category == "anonymous_target"
|
|
assert TorchLumaBlobDetector.metadata.weights == "none"
|