| | |
| | from __future__ import division |
| |
|
| | import numpy as np |
| |
|
| | from annotator.uniformer.mmcv.image import rgb2bgr |
| | from annotator.uniformer.mmcv.video import flowread |
| | from .image import imshow |
| |
|
| |
|
| | def flowshow(flow, win_name='', wait_time=0): |
| | """Show optical flow. |
| | |
| | Args: |
| | flow (ndarray or str): The optical flow to be displayed. |
| | win_name (str): The window name. |
| | wait_time (int): Value of waitKey param. |
| | """ |
| | flow = flowread(flow) |
| | flow_img = flow2rgb(flow) |
| | imshow(rgb2bgr(flow_img), win_name, wait_time) |
| |
|
| |
|
| | def flow2rgb(flow, color_wheel=None, unknown_thr=1e6): |
| | """Convert flow map to RGB image. |
| | |
| | Args: |
| | flow (ndarray): Array of optical flow. |
| | color_wheel (ndarray or None): Color wheel used to map flow field to |
| | RGB colorspace. Default color wheel will be used if not specified. |
| | 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 |
| | if color_wheel is None: |
| | color_wheel = 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 |
| |
|
| | rad = np.sqrt(dx**2 + dy**2) |
| | if np.any(rad > np.finfo(float).eps): |
| | max_rad = np.max(rad) |
| | dx /= max_rad |
| | dy /= max_rad |
| |
|
| | rad = np.sqrt(dx**2 + dy**2) |
| | angle = np.arctan2(-dy, -dx) / np.pi |
| |
|
| | bin_real = (angle + 1) / 2 * (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, :] |
| | small_ind = rad <= 1 |
| | flow_img[small_ind] = 1 - rad[small_ind, None] * (1 - flow_img[small_ind]) |
| | flow_img[np.logical_not(small_ind)] *= 0.75 |
| |
|
| | flow_img[ignore_inds, :] = 0 |
| |
|
| | return flow_img |
| |
|
| |
|
| | def make_color_wheel(bins=None): |
| | """Build a color wheel. |
| | |
| | Args: |
| | bins(list or tuple, optional): Specify the number of bins for each |
| | color range, corresponding to six ranges: red -> yellow, |
| | yellow -> green, green -> cyan, cyan -> blue, blue -> magenta, |
| | magenta -> red. [15, 6, 4, 11, 13, 6] is used for default |
| | (see Middlebury). |
| | |
| | Returns: |
| | ndarray: Color wheel of shape (total_bins, 3). |
| | """ |
| | if bins is None: |
| | bins = [15, 6, 4, 11, 13, 6] |
| | assert len(bins) == 6 |
| |
|
| | RY, YG, GC, CB, BM, MR = tuple(bins) |
| |
|
| | ry = [1, np.arange(RY) / RY, 0] |
| | yg = [1 - np.arange(YG) / YG, 1, 0] |
| | gc = [0, 1, np.arange(GC) / GC] |
| | cb = [0, 1 - np.arange(CB) / CB, 1] |
| | bm = [np.arange(BM) / BM, 0, 1] |
| | mr = [1, 0, 1 - np.arange(MR) / MR] |
| |
|
| | num_bins = RY + YG + GC + CB + BM + MR |
| |
|
| | color_wheel = np.zeros((3, num_bins), dtype=np.float32) |
| |
|
| | col = 0 |
| | for i, color in enumerate([ry, yg, gc, cb, bm, mr]): |
| | for j in range(3): |
| | color_wheel[j, col:col + bins[i]] = color[j] |
| | col += bins[i] |
| |
|
| | return color_wheel.T |
| |
|
