feat: add deterministic marker-based size bucketing
This commit is contained in:
parent
9af936518d
commit
b9ddbee81e
2 changed files with 86 additions and 0 deletions
39
bookmark/sizing.py
Normal file
39
bookmark/sizing.py
Normal file
|
|
@ -0,0 +1,39 @@
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class MarkerRegion:
|
||||||
|
name: str
|
||||||
|
bucket: str
|
||||||
|
x: int
|
||||||
|
y: int
|
||||||
|
width: int
|
||||||
|
height: int
|
||||||
|
|
||||||
|
|
||||||
|
def _mean_region_diff(frame: np.ndarray, reference_frame: np.ndarray, marker: MarkerRegion) -> float:
|
||||||
|
frame_region = frame[marker.y:marker.y + marker.height, marker.x:marker.x + marker.width]
|
||||||
|
reference_region = reference_frame[marker.y:marker.y + marker.height, marker.x:marker.x + marker.width]
|
||||||
|
diff = np.abs(frame_region.astype(int) - reference_region.astype(int))
|
||||||
|
return float(diff.mean())
|
||||||
|
|
||||||
|
|
||||||
|
def determine_size_bucket(
|
||||||
|
frame: np.ndarray,
|
||||||
|
reference_frame: np.ndarray,
|
||||||
|
markers: list[MarkerRegion],
|
||||||
|
threshold: float = 30.0,
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
`markers` must be ordered smallest -> largest bucket boundary. Returns
|
||||||
|
the bucket of the largest boundary whose marker region is occluded
|
||||||
|
(covered by the book) relative to `reference_frame`. Defaults to the
|
||||||
|
smallest bucket if no marker is occluded.
|
||||||
|
"""
|
||||||
|
bucket = markers[0].bucket
|
||||||
|
for marker in markers:
|
||||||
|
if _mean_region_diff(frame, reference_frame, marker) > threshold:
|
||||||
|
bucket = marker.bucket
|
||||||
|
return bucket
|
||||||
47
tests/test_sizing.py
Normal file
47
tests/test_sizing.py
Normal file
|
|
@ -0,0 +1,47 @@
|
||||||
|
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"
|
||||||
Loading…
Reference in a new issue