bookmark/tests/test_stable_centered_strategy.py

53 lines
1.9 KiB
Python

from bookmark.capture_strategies.stable_centered import StableCenteredStrategy
import numpy as np
def make_reference(size=100, color=200):
return np.full((size, size, 3), color, dtype=np.uint8)
def make_book_frame(size=100, ref_color=200, book_color=50, box=(20, 20, 60, 60)):
frame = np.full((size, size, 3), ref_color, dtype=np.uint8)
x, y, w, h = box
frame[y:y + h, x:x + w] = book_color
return frame
def test_no_fire_on_empty_reference_frame():
reference = make_reference()
strategy = StableCenteredStrategy(reference_frame=reference, guide_box=(20, 20, 60, 60), stable_frames=3)
assert strategy.observe(reference) is False
def test_no_fire_when_book_too_small():
reference = make_reference()
small_book = make_book_frame(box=(40, 40, 10, 10))
strategy = StableCenteredStrategy(
reference_frame=reference, guide_box=(20, 20, 60, 60), min_fill_ratio=0.5, stable_frames=3
)
for _ in range(5):
assert strategy.observe(small_book) is False
def test_fires_after_stable_frames_with_book_filling_guide_box():
reference = make_reference()
book_frame = make_book_frame(box=(20, 20, 60, 60))
strategy = StableCenteredStrategy(
reference_frame=reference, guide_box=(20, 20, 60, 60), min_fill_ratio=0.5, stable_frames=3
)
assert strategy.observe(book_frame) is False
assert strategy.observe(book_frame) is False
assert strategy.observe(book_frame) is False
assert strategy.observe(book_frame) is True
def test_fires_only_once():
reference = make_reference()
book_frame = make_book_frame(box=(20, 20, 60, 60))
strategy = StableCenteredStrategy(
reference_frame=reference, guide_box=(20, 20, 60, 60), min_fill_ratio=0.5, stable_frames=2
)
strategy.observe(book_frame)
strategy.observe(book_frame)
assert strategy.observe(book_frame) is True
assert strategy.observe(book_frame) is False