ykcheungab
116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
"""orisys.util.draw_arrows: mask handling and the out-of-frame tip regression."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from orisys.util import _prepare_draw_mask, draw_arrows
|
|
|
|
|
|
def flow_of(shape, dx, dy):
|
|
flow = np.zeros((*shape, 2), dtype=np.float32)
|
|
flow[..., 0] = dx / 2.0 # draw_arrows doubles the stored value
|
|
flow[..., 1] = dy / 2.0
|
|
return flow
|
|
|
|
|
|
def background(frame):
|
|
return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
|
|
|
|
|
|
def test_out_of_frame_arrow_tip_is_dropped_not_raised():
|
|
"""Regression: mask[end_y, end_x] was indexed before the bounds check.
|
|
|
|
A 40 px grid with the default edge_margin of 15 lets an arrow tip leave the
|
|
frame, which used to raise IndexError as soon as a valid_mask was supplied.
|
|
"""
|
|
frame = np.full((48, 48), 100, dtype=np.uint8)
|
|
flow = flow_of((48, 48), dx=40.0, dy=0.0)
|
|
mask = np.ones((48, 48), dtype=bool)
|
|
|
|
out = draw_arrows(frame, flow, grid_spacing=40, margin=0.8, valid_mask=mask)
|
|
|
|
assert out.shape == (48, 48, 3)
|
|
# The only grid point (20, 20) has a tip at x=58 in a 48 px frame, so the
|
|
# bounds check must have discarded it: nothing was drawn.
|
|
assert np.array_equal(out, background(frame))
|
|
|
|
|
|
def test_out_of_frame_tip_matches_render_without_mask():
|
|
"""Clamping the mask lookup must not change which points are valid."""
|
|
frame = np.full((48, 48), 100, dtype=np.uint8)
|
|
flow = flow_of((48, 48), dx=40.0, dy=0.0)
|
|
mask = np.ones((48, 48), dtype=bool)
|
|
|
|
with_mask = draw_arrows(frame, flow, grid_spacing=40, valid_mask=mask)
|
|
without = draw_arrows(frame, flow, grid_spacing=40)
|
|
assert np.array_equal(with_mask, without)
|
|
|
|
|
|
def test_in_frame_arrows_are_drawn_through_an_all_true_mask():
|
|
frame = np.full((96, 96), 90, dtype=np.uint8)
|
|
flow = flow_of((96, 96), dx=6.0, dy=3.0)
|
|
mask = np.ones((96, 96), dtype=bool)
|
|
|
|
out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0,
|
|
threshold=0.1, valid_mask=mask)
|
|
assert not np.array_equal(out, background(frame))
|
|
|
|
|
|
def test_all_false_mask_draws_nothing():
|
|
frame = np.full((96, 96), 90, dtype=np.uint8)
|
|
flow = flow_of((96, 96), dx=6.0, dy=3.0)
|
|
mask = np.zeros((96, 96), dtype=bool)
|
|
|
|
out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0,
|
|
threshold=0.1, valid_mask=mask)
|
|
assert np.array_equal(out, background(frame))
|
|
|
|
|
|
def test_mask_can_exclude_one_side():
|
|
"""Endpoints outside the mask are skipped even when inside the frame."""
|
|
frame = np.full((96, 96), 90, dtype=np.uint8)
|
|
flow = flow_of((96, 96), dx=8.0, dy=0.0)
|
|
mask = np.zeros((96, 96), dtype=bool)
|
|
mask[24:48, 24:72] = True # only the upper middle band
|
|
|
|
out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0,
|
|
threshold=0.1, valid_mask=mask)
|
|
assert not np.array_equal(out, background(frame))
|
|
# Rows below the band stay untouched.
|
|
assert np.array_equal(out[80:, :, :], background(frame)[80:, :, :])
|
|
|
|
|
|
def test_below_threshold_points_get_the_dot_colour():
|
|
frame = np.full((96, 96), 90, dtype=np.uint8)
|
|
flow = flow_of((96, 96), dx=0.05, dy=0.05)
|
|
out = draw_arrows(frame, flow, grid_spacing=12, arrow_scale=1.0,
|
|
threshold=1.0, below_threshold_color=(0, 0, 255))
|
|
# OpenCV colours are BGR, so (0, 0, 255) is red.
|
|
red = (out[:, :, 2] > 150) & (out[:, :, 0] < 100) & (out[:, :, 1] < 100)
|
|
assert red.sum() > 0
|
|
|
|
|
|
def test_prepare_draw_mask_resizes_to_frame():
|
|
mask = np.zeros((10, 20), dtype=bool)
|
|
mask[:, :10] = True
|
|
out = _prepare_draw_mask(mask, (40, 40))
|
|
assert out.shape == (40, 40)
|
|
assert out[:, :20].all()
|
|
assert not out[:, 20:].any()
|
|
|
|
|
|
def test_prepare_draw_mask_passes_none_through():
|
|
assert _prepare_draw_mask(None, (10, 10)) is None
|
|
|
|
|
|
def test_non_square_grid_does_not_raise_with_mask():
|
|
"""Height and width differ, so rr/cc broadcasting must stay (N, N)."""
|
|
frame = np.full((60, 120), 80, dtype=np.uint8)
|
|
flow = flow_of((60, 120), dx=5.0, dy=5.0)
|
|
out = draw_arrows(frame, flow, grid_spacing=10, arrow_scale=1.0,
|
|
threshold=0.1, valid_mask=np.ones((60, 120), dtype=bool))
|
|
assert out.shape == (60, 120, 3)
|
|
assert not np.array_equal(out, background(frame))
|