bookmark/tests/test_sizing.py

47 lines
1.5 KiB
Python

import numpy as np
from bookmark.sizing import MarkerRegion, determine_size_bucket
def make_reference(size=100, color=200):
return np.full((size, size, 3), color, dtype=np.uint8)
MARKERS = [
MarkerRegion(name="small_boundary", bucket="small", x=0, y=0, width=10, height=10),
MarkerRegion(name="medium_boundary", bucket="medium", x=50, y=0, width=10, height=10),
MarkerRegion(name="large_boundary", bucket="large", x=0, y=50, width=10, height=10),
]
def occlude(frame, marker, color=50):
frame = frame.copy()
frame[marker.y:marker.y + marker.height, marker.x:marker.x + marker.width] = color
return frame
def test_no_markers_occluded_returns_smallest_bucket():
reference = make_reference()
frame = reference.copy()
assert determine_size_bucket(frame, reference, MARKERS) == "small"
def test_only_small_marker_occluded_returns_small():
reference = make_reference()
frame = occlude(reference, MARKERS[0])
assert determine_size_bucket(frame, reference, MARKERS) == "small"
def test_small_and_medium_occluded_returns_medium():
reference = make_reference()
frame = occlude(reference, MARKERS[0])
frame = occlude(frame, MARKERS[1])
assert determine_size_bucket(frame, reference, MARKERS) == "medium"
def test_all_markers_occluded_returns_large():
reference = make_reference()
frame = reference.copy()
for marker in MARKERS:
frame = occlude(frame, marker)
assert determine_size_bucket(frame, reference, MARKERS) == "large"