def _flow_to_direction_and_magnitude(flow, unknown_thr=1e6):
"""Convert flow map to RGB image.
Args:
flow (ndarray): Array of optical flow.
unknown_thr (str): Values above this threshold will be marked as
unknown and thus ignored.
Returns:
ndarray: RGB image that can be visualized.
"""
assert flow.ndim == 3 and flow.shape[-1] == 2
color_wheel = mmcv.make_color_wheel()
assert color_wheel.ndim == 2 and color_wheel.shape[1] == 3
num_bins = color_wheel.shape[0]
dx = flow[:, :, 0].copy()
dy = flow[:, :, 1].copy()
ignore_inds = (np.isnan(dx) | np.isnan(dy) | (np.abs(dx) > unknown_thr) |
(np.abs(dy) > unknown_thr))
dx[ignore_inds] = 0
dy[ignore_inds] = 0
flow_magnitude = np.sqrt(dx**2 + dy**2)
if np.any(flow_magnitude > np.finfo(float).eps):
max_rad = np.max(flow_magnitude)
dx /= max_rad
dy /= max_rad
flow_magnitude = np.sqrt(dx**2 + dy**2)
flow_direction = np.arctan2(-dy, -dx) / np.pi # -1,1
bin_real = (flow_direction + 1) / 2 * (num_bins - 1) # [0,num_bins-1)
bin_left = np.floor(bin_real).astype(int)
bin_right = (bin_left + 1) % num_bins
w = (bin_real - bin_left.astype(np.float32))[..., None]
flow_img = (1 - w) * color_wheel[bin_left, :] + \
w * color_wheel[bin_right, :]
direction_map = flow_img.copy()
small_ind = flow_magnitude <= 1
flow_img[small_ind] = 1 - flow_magnitude[small_ind, None] * \
(1 - flow_img[small_ind])
flow_img[np.logical_not(small_ind)] *= 0.75
flow_img[ignore_inds, :] = 0
return dict(flow=flow_img,
direction=direction_map,
magnitude=flow_magnitude)
def test_make_color_wheel():
default_color_wheel = mmcv.make_color_wheel()
color_wheel = mmcv.make_color_wheel([2, 2, 2, 2, 2, 2])
# yapf: disable
assert_array_equal(default_color_wheel, np.array(
[[1. , 0. , 0. ],
[1. , 0.06666667, 0. ],
[1. , 0.13333334, 0. ],
[1. , 0.2 , 0. ],
[1. , 0.26666668, 0. ],
[1. , 0.33333334, 0. ],
[1. , 0.4 , 0. ],
[1. , 0.46666667, 0. ],
[1. , 0.53333336, 0. ],
[1. , 0.6 , 0. ],
[1. , 0.6666667 , 0. ],
[1. , 0.73333335, 0. ],
[1. , 0.8 , 0. ],
[1. , 0.8666667 , 0. ],
[1. , 0.93333334, 0. ],
[1. , 1. , 0. ],
[0.8333333 , 1. , 0. ],
[0.6666667 , 1. , 0. ],
[0.5 , 1. , 0. ],
[0.33333334, 1. , 0. ],
[0.16666667, 1. , 0. ],
[0. , 1. , 0. ],
[0. , 1. , 0.25 ],
[0. , 1. , 0.5 ],
[0. , 1. , 0.75 ],
[0. , 1. , 1. ],
[0. , 0.90909094, 1. ],
[0. , 0.8181818 , 1. ],
[0. , 0.72727275, 1. ],
[0. , 0.6363636 , 1. ],
[0. , 0.54545456, 1. ],
[0. , 0.45454547, 1. ],
[0. , 0.36363637, 1. ],
[0. , 0.27272728, 1. ],
[0. , 0.18181819, 1. ],
[0. , 0.09090909, 1. ],
[0. , 0. , 1. ],
[0.07692308, 0. , 1. ],
[0.15384616, 0. , 1. ],
[0.23076923, 0. , 1. ],
[0.30769232, 0. , 1. ],
[0.3846154 , 0. , 1. ],
[0.46153846, 0. , 1. ],
[0.53846157, 0. , 1. ],
[0.61538464, 0. , 1. ],
[0.6923077 , 0. , 1. ],
[0.7692308 , 0. , 1. ],
[0.84615386, 0. , 1. ],
[0.9230769 , 0. , 1. ],
[1. , 0. , 1. ],
[1. , 0. , 0.8333333 ],
[1. , 0. , 0.6666667 ],
[1. , 0. , 0.5 ],
[1. , 0. , 0.33333334],
[1. , 0. , 0.16666667]], dtype=np.float32))
assert_array_equal(
color_wheel,
np.array([[1., 0. , 0. ],
[1. , 0.5, 0. ],
[1. , 1. , 0. ],
[0.5, 1. , 0. ],
[0. , 1. , 0. ],
[0. , 1. , 0.5],
[0. , 1. , 1. ],
[0. , 0.5, 1. ],
[0. , 0. , 1. ],
[0.5, 0. , 1. ],
[1. , 0. , 1. ],
[1. , 0. , 0.5]], dtype=np.float32))