Spaces:
Build error
Build error
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from typing import Optional, Union | |
| import numpy as np | |
| from mmcv.image import rgb2bgr | |
| from mmcv.video import flowread | |
| from .image import imshow | |
| def flowshow(flow: Union[np.ndarray, str], | |
| win_name: str = '', | |
| wait_time: int = 0) -> None: | |
| """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: np.ndarray, | |
| color_wheel: Optional[np.ndarray] = None, | |
| unknown_thr: float = 1e6) -> np.ndarray: | |
| """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 (float): 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: Optional[Union[list, tuple]] = None) -> np.ndarray: | |
| """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 | |