Unverified Commit 882cab77 authored by Song Lin's avatar Song Lin Committed by GitHub
Browse files

[Enhancement] Add type annotations for mmcv/visualization (#1947)



* Add type hints

* Replace mmcv is_str by isinstance str

* Add type hints in optflow.py

* Update mmcv/visualization/optflow.py
Co-authored-by: default avatarZaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update mmcv/visualization/color.py
Co-authored-by: default avatarZaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Fix tuple hints
Co-authored-by: default avatarZaida Zhou <58739961+zhouzaida@users.noreply.github.com>
parent 105fe550
# Copyright (c) OpenMMLab. All rights reserved. # Copyright (c) OpenMMLab. All rights reserved.
from enum import Enum from enum import Enum
from typing import Union
import numpy as np import numpy as np
...@@ -21,7 +22,7 @@ class Color(Enum): ...@@ -21,7 +22,7 @@ class Color(Enum):
black = (0, 0, 0) black = (0, 0, 0)
def color_val(color): def color_val(color: Union[Color, str, tuple, int, np.ndarray]) -> tuple:
"""Convert various input to color tuples. """Convert various input to color tuples.
Args: Args:
...@@ -31,7 +32,7 @@ def color_val(color): ...@@ -31,7 +32,7 @@ def color_val(color):
tuple[int]: A tuple of 3 integers indicating BGR channels. tuple[int]: A tuple of 3 integers indicating BGR channels.
""" """
if is_str(color): if is_str(color):
return Color[color].value return Color[color].value # type: ignore
elif isinstance(color, Color): elif isinstance(color, Color):
return color.value return color.value
elif isinstance(color, tuple): elif isinstance(color, tuple):
......
# Copyright (c) OpenMMLab. All rights reserved. # Copyright (c) OpenMMLab. All rights reserved.
from __future__ import division from __future__ import division
from typing import Optional, Union
import numpy as np import numpy as np
...@@ -8,7 +9,9 @@ from mmcv.video import flowread ...@@ -8,7 +9,9 @@ from mmcv.video import flowread
from .image import imshow from .image import imshow
def flowshow(flow, win_name='', wait_time=0): def flowshow(flow: Union[np.ndarray, str],
win_name: str = '',
wait_time: int = 0) -> None:
"""Show optical flow. """Show optical flow.
Args: Args:
...@@ -21,14 +24,16 @@ def flowshow(flow, win_name='', wait_time=0): ...@@ -21,14 +24,16 @@ def flowshow(flow, win_name='', wait_time=0):
imshow(rgb2bgr(flow_img), win_name, wait_time) imshow(rgb2bgr(flow_img), win_name, wait_time)
def flow2rgb(flow, color_wheel=None, unknown_thr=1e6): def flow2rgb(flow: np.ndarray,
color_wheel: Optional[np.ndarray] = None,
unknown_thr: float = 1e6) -> np.ndarray:
"""Convert flow map to RGB image. """Convert flow map to RGB image.
Args: Args:
flow (ndarray): Array of optical flow. flow (ndarray): Array of optical flow.
color_wheel (ndarray or None): Color wheel used to map flow field to color_wheel (ndarray or None): Color wheel used to map flow field to
RGB colorspace. Default color wheel will be used if not specified. RGB colorspace. Default color wheel will be used if not specified.
unknown_thr (str): Values above this threshold will be marked as unknown_thr (float): Values above this threshold will be marked as
unknown and thus ignored. unknown and thus ignored.
Returns: Returns:
...@@ -73,7 +78,7 @@ def flow2rgb(flow, color_wheel=None, unknown_thr=1e6): ...@@ -73,7 +78,7 @@ def flow2rgb(flow, color_wheel=None, unknown_thr=1e6):
return flow_img return flow_img
def make_color_wheel(bins=None): def make_color_wheel(bins: Optional[Union[list, tuple]] = None) -> np.ndarray:
"""Build a color wheel. """Build a color wheel.
Args: Args:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment