mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Add type hints for mmcv/image (#2089) 

* Fix typehint

* minor fix

* minor fix

* minor fix

Co-authored-by: zhouzaida <zhouzaida@163.com>
pull/2349/head
tripleMu 2022-09-12 17:16:55 +08:00 committed by Zaida Zhou
parent f7c01d44b1
commit 0412097ec2
3 changed files with 114 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

143
mmcv/image/geometric.py
View File

@ -1,7 +1,7 @@
# Copyright (c) OpenMMLab. All rights reserved. # Copyright (c) OpenMMLab. All rights reserved.
import numbers import numbers
import warnings import warnings
from typing import Optional, Tuple from typing import List, Optional, Tuple, Union, no_type_check
import cv2 import cv2
import numpy as np import numpy as np
@ -15,7 +15,10 @@ except ImportError:
Image = None Image = None
def _scale_size(size, scale): def _scale_size(
size: Tuple[int, int],
scale: Union[float, int, tuple],
) -> Tuple[int, int]:
"""Rescale a size by a ratio. """Rescale a size by a ratio.
Args: Args:
@ -72,12 +75,14 @@ if Image is not None:
} }
def imresize(img, def imresize(
size, img: np.ndarray,
return_scale=False, size: Tuple[int, int],
interpolation='bilinear', return_scale: bool = False,
out=None, interpolation: str = 'bilinear',
backend=None): out: Optional[np.ndarray] = None,
backend: Optional[str] = None
) -> Union[Tuple[np.ndarray, float, float], np.ndarray]:
"""Resize image to a given size. """Resize image to a given size.
Args: Args:
@ -119,15 +124,18 @@ def imresize(img,
return resized_img, w_scale, h_scale return resized_img, w_scale, h_scale
def imresize_to_multiple(img, @no_type_check
divisor, def imresize_to_multiple(
size=None, img: np.ndarray,
scale_factor=None, divisor: Union[int, Tuple[int, int]],
keep_ratio=False, size: Union[int, Tuple[int, int], None] = None,
return_scale=False, scale_factor: Union[float, Tuple[float, float], None] = None,
interpolation='bilinear', keep_ratio: bool = False,
out=None, return_scale: bool = False,
backend=None): interpolation: str = 'bilinear',
out: Optional[np.ndarray] = None,
backend: Optional[str] = None
) -> Union[Tuple[np.ndarray, float, float], np.ndarray]:
"""Resize image according to a given size or scale factor and then rounds """Resize image according to a given size or scale factor and then rounds
up the the resized or rescaled image size to the nearest value that can be up the the resized or rescaled image size to the nearest value that can be
divided by the divisor. divided by the divisor.
@ -183,11 +191,13 @@ def imresize_to_multiple(img,
return resized_img return resized_img
def imresize_like(img, def imresize_like(
dst_img, img: np.ndarray,
return_scale=False, dst_img: np.ndarray,
interpolation='bilinear', return_scale: bool = False,
backend=None): interpolation: str = 'bilinear',
backend: Optional[str] = None
) -> Union[Tuple[np.ndarray, float, float], np.ndarray]:
"""Resize image to the same size of a given image. """Resize image to the same size of a given image.
Args: Args:
@ -205,7 +215,9 @@ def imresize_like(img,
return imresize(img, (w, h), return_scale, interpolation, backend=backend) return imresize(img, (w, h), return_scale, interpolation, backend=backend)
def rescale_size(old_size, scale, return_scale=False): def rescale_size(old_size: tuple,
scale: Union[float, int, tuple],
return_scale: bool = False) -> tuple:
"""Calculate the new size to be rescaled to. """Calculate the new size to be rescaled to.
Args: Args:
@ -242,11 +254,13 @@ def rescale_size(old_size, scale, return_scale=False):
return new_size return new_size
def imrescale(img, def imrescale(
scale, img: np.ndarray,
return_scale=False, scale: Union[float, Tuple[int, int]],
interpolation='bilinear', return_scale: bool = False,
backend=None): interpolation: str = 'bilinear',
backend: Optional[str] = None
) -> Union[np.ndarray, Tuple[np.ndarray, float]]:
"""Resize image while keeping the aspect ratio. """Resize image while keeping the aspect ratio.
Args: Args:
@ -273,7 +287,7 @@ def imrescale(img,
return rescaled_img return rescaled_img
def imflip(img, direction='horizontal'): def imflip(img: np.ndarray, direction: str = 'horizontal') -> np.ndarray:
"""Flip an image horizontally or vertically. """Flip an image horizontally or vertically.
Args: Args:
@ -293,7 +307,7 @@ def imflip(img, direction='horizontal'):
return np.flip(img, axis=(0, 1)) return np.flip(img, axis=(0, 1))
def imflip_(img, direction='horizontal'): def imflip_(img: np.ndarray, direction: str = 'horizontal') -> np.ndarray:
"""Inplace flip an image horizontally or vertically. """Inplace flip an image horizontally or vertically.
Args: Args:
@ -373,7 +387,7 @@ def imrotate(img: np.ndarray,
return rotated return rotated
def bbox_clip(bboxes, img_shape): def bbox_clip(bboxes: np.ndarray, img_shape: Tuple[int, int]) -> np.ndarray:
"""Clip bboxes to fit the image shape. """Clip bboxes to fit the image shape.
Args: Args:
@ -391,7 +405,9 @@ def bbox_clip(bboxes, img_shape):
return clipped_bboxes return clipped_bboxes
def bbox_scaling(bboxes, scale, clip_shape=None): def bbox_scaling(bboxes: np.ndarray,
scale: float,
clip_shape: Optional[Tuple[int, int]] = None) -> np.ndarray:
"""Scaling bboxes w.r.t the box center. """Scaling bboxes w.r.t the box center.
Args: Args:
@ -417,7 +433,12 @@ def bbox_scaling(bboxes, scale, clip_shape=None):
return scaled_bboxes return scaled_bboxes
def imcrop(img, bboxes, scale=1.0, pad_fill=None): def imcrop(
img: np.ndarray,
bboxes: np.ndarray,
scale: float = 1.0,
pad_fill: Union[float, list, None] = None
) -> Union[np.ndarray, List[np.ndarray]]:
"""Crop image patches. """Crop image patches.
3 steps: scale the bboxes -> clip bboxes -> crop and pad. 3 steps: scale the bboxes -> clip bboxes -> crop and pad.
@ -450,10 +471,12 @@ def imcrop(img, bboxes, scale=1.0, pad_fill=None):
patch = img[y1:y2 + 1, x1:x2 + 1, ...] patch = img[y1:y2 + 1, x1:x2 + 1, ...]
else: else:
_x1, _y1, _x2, _y2 = tuple(scaled_bboxes[i, :]) _x1, _y1, _x2, _y2 = tuple(scaled_bboxes[i, :])
patch_h = _y2 - _y1 + 1
patch_w = _x2 - _x1 + 1
if chn == 1: if chn == 1:
patch_shape = (_y2 - _y1 + 1, _x2 - _x1 + 1) patch_shape = (patch_h, patch_w)
else: else:
patch_shape = (_y2 - _y1 + 1, _x2 - _x1 + 1, chn) patch_shape = (patch_h, patch_w, chn) # type: ignore
patch = np.array( patch = np.array(
pad_fill, dtype=img.dtype) * np.ones( pad_fill, dtype=img.dtype) * np.ones(
patch_shape, dtype=img.dtype) patch_shape, dtype=img.dtype)
@ -471,12 +494,12 @@ def imcrop(img, bboxes, scale=1.0, pad_fill=None):
return patches return patches
def impad(img, def impad(img: np.ndarray,
*, *,
shape=None, shape: Optional[Tuple[int, int]] = None,
padding=None, padding: Union[int, tuple, None] = None,
pad_val=0, pad_val: Union[float, List] = 0,
padding_mode='constant'): padding_mode: str = 'constant') -> np.ndarray:
"""Pad the given image to a certain shape or pad on all sides with """Pad the given image to a certain shape or pad on all sides with
specified padding mode and padding value. specified padding mode and padding value.
@ -555,7 +578,9 @@ def impad(img,
return img return img
def impad_to_multiple(img, divisor, pad_val=0): def impad_to_multiple(img: np.ndarray,
divisor: int,
pad_val: Union[float, List] = 0) -> np.ndarray:
"""Pad an image to ensure each edge to be multiple to some number. """Pad an image to ensure each edge to be multiple to some number.
Args: Args:
@ -571,7 +596,9 @@ def impad_to_multiple(img, divisor, pad_val=0):
return impad(img, shape=(pad_h, pad_w), pad_val=pad_val) return impad(img, shape=(pad_h, pad_w), pad_val=pad_val)
def cutout(img, shape, pad_val=0): def cutout(img: np.ndarray,
shape: Union[int, Tuple[int, int]],
pad_val: Union[int, float, tuple] = 0) -> np.ndarray:
"""Randomly cut out a rectangle from the original img. """Randomly cut out a rectangle from the original img.
Args: Args:
@ -615,7 +642,7 @@ def cutout(img, shape, pad_val=0):
if img.ndim == 2: if img.ndim == 2:
patch_shape = (y2 - y1, x2 - x1) patch_shape = (y2 - y1, x2 - x1)
else: else:
patch_shape = (y2 - y1, x2 - x1, channels) patch_shape = (y2 - y1, x2 - x1, channels) # type: ignore
img_cutout = img.copy() img_cutout = img.copy()
patch = np.array( patch = np.array(
@ -626,7 +653,8 @@ def cutout(img, shape, pad_val=0):
return img_cutout return img_cutout
def _get_shear_matrix(magnitude, direction='horizontal'): def _get_shear_matrix(magnitude: Union[int, float],
direction: str = 'horizontal') -> np.ndarray:
"""Generate the shear matrix for transformation. """Generate the shear matrix for transformation.
Args: Args:
@ -644,11 +672,11 @@ def _get_shear_matrix(magnitude, direction='horizontal'):
return shear_matrix return shear_matrix
def imshear(img, def imshear(img: np.ndarray,
magnitude, magnitude: Union[int, float],
direction='horizontal', direction: str = 'horizontal',
border_value=0, border_value: Union[int, Tuple[int, int]] = 0,
interpolation='bilinear'): interpolation: str = 'bilinear') -> np.ndarray:
"""Shear an image. """Shear an image.
Args: Args:
@ -672,7 +700,7 @@ def imshear(img,
elif img.ndim == 3: elif img.ndim == 3:
channels = img.shape[-1] channels = img.shape[-1]
if isinstance(border_value, int): if isinstance(border_value, int):
border_value = tuple([border_value] * channels) border_value = tuple([border_value] * channels) # type: ignore
elif isinstance(border_value, tuple): elif isinstance(border_value, tuple):
assert len(border_value) == channels, \ assert len(border_value) == channels, \
'Expected the num of elements in tuple equals the channels' \ 'Expected the num of elements in tuple equals the channels' \
@ -690,12 +718,13 @@ def imshear(img,
# greater than 3 (e.g. shearing masks whose channels large # greater than 3 (e.g. shearing masks whose channels large
# than 3) will raise TypeError in `cv2.warpAffine`. # than 3) will raise TypeError in `cv2.warpAffine`.
# Here simply slice the first 3 values in `border_value`. # Here simply slice the first 3 values in `border_value`.
borderValue=border_value[:3], borderValue=border_value[:3], # type: ignore
flags=cv2_interp_codes[interpolation]) flags=cv2_interp_codes[interpolation])
return sheared return sheared
def _get_translate_matrix(offset, direction='horizontal'): def _get_translate_matrix(offset: Union[int, float],
direction: str = 'horizontal') -> np.ndarray:
"""Generate the translate matrix. """Generate the translate matrix.
Args: Args:
@ -713,11 +742,11 @@ def _get_translate_matrix(offset, direction='horizontal'):
return translate_matrix return translate_matrix
def imtranslate(img, def imtranslate(img: np.ndarray,
offset, offset: Union[int, float],
direction='horizontal', direction: str = 'horizontal',
border_value=0, border_value: Union[int, tuple] = 0,
interpolation='bilinear'): interpolation: str = 'bilinear') -> np.ndarray:
"""Translate an image. """Translate an image.
Args: Args:

37
mmcv/image/io.py
View File

@ -3,6 +3,7 @@ import io
import os.path as osp import os.path as osp
import warnings import warnings
from pathlib import Path from pathlib import Path
from typing import Optional, Union
import cv2 import cv2
import numpy as np import numpy as np
@ -41,7 +42,7 @@ imread_flags = {
imread_backend = 'cv2' imread_backend = 'cv2'
def use_backend(backend): def use_backend(backend: str) -> None:
"""Select a backend for image decoding. """Select a backend for image decoding.
Args: Args:
@ -67,7 +68,7 @@ def use_backend(backend):
raise ImportError('`tifffile` is not installed') raise ImportError('`tifffile` is not installed')
def _jpegflag(flag='color', channel_order='bgr'): def _jpegflag(flag: str = 'color', channel_order: str = 'bgr'):
channel_order = channel_order.lower() channel_order = channel_order.lower()
if channel_order not in ['rgb', 'bgr']: if channel_order not in ['rgb', 'bgr']:
raise ValueError('channel order must be either "rgb" or "bgr"') raise ValueError('channel order must be either "rgb" or "bgr"')
@ -83,7 +84,9 @@ def _jpegflag(flag='color', channel_order='bgr'):
raise ValueError('flag must be "color" or "grayscale"') raise ValueError('flag must be "color" or "grayscale"')
def _pillow2array(img, flag='color', channel_order='bgr'): def _pillow2array(img,
flag: str = 'color',
channel_order: str = 'bgr') -> np.ndarray:
"""Convert a pillow image to numpy array. """Convert a pillow image to numpy array.
Args: Args:
@ -138,11 +141,11 @@ def _pillow2array(img, flag='color', channel_order='bgr'):
return array return array
def imread(img_or_path, def imread(img_or_path: Union[np.ndarray, str, Path],
flag='color', flag: str = 'color',
channel_order='bgr', channel_order: str = 'bgr',
backend=None, backend: Optional[str] = None,
file_client_args=None): file_client_args: Optional[dict] = None) -> np.ndarray:
"""Read an image. """Read an image.
Note: Note:
@ -206,7 +209,10 @@ def imread(img_or_path,
'a pathlib.Path object') 'a pathlib.Path object')
def imfrombytes(content, flag='color', channel_order='bgr', backend=None): def imfrombytes(content: bytes,
flag: str = 'color',
channel_order: str = 'bgr',
backend: Optional[str] = None) -> np.ndarray:
"""Read an image from bytes. """Read an image from bytes.
Args: Args:
@ -239,7 +245,8 @@ def imfrombytes(content, flag='color', channel_order='bgr', backend=None):
f'backend: {backend} is not supported. Supported ' f'backend: {backend} is not supported. Supported '
"backends are 'cv2', 'turbojpeg', 'pillow', 'tifffile'") "backends are 'cv2', 'turbojpeg', 'pillow', 'tifffile'")
if backend == 'turbojpeg': if backend == 'turbojpeg':
img = jpeg.decode(content, _jpegflag(flag, channel_order)) img = jpeg.decode( # type: ignore
content, _jpegflag(flag, channel_order))
if img.shape[-1] == 1: if img.shape[-1] == 1:
img = img[:, :, 0] img = img[:, :, 0]
return img return img
@ -261,11 +268,11 @@ def imfrombytes(content, flag='color', channel_order='bgr', backend=None):
return img return img
def imwrite(img, def imwrite(img: np.ndarray,
file_path, file_path: str,
params=None, params: Optional[list] = None,
auto_mkdir=None, auto_mkdir: Optional[bool] = None,
file_client_args=None): file_client_args: Optional[dict] = None) -> bool:
"""Write image to file. """Write image to file.
Note: Note:

7
mmcv/image/misc.py
View File

@ -1,4 +1,6 @@
# Copyright (c) OpenMMLab. All rights reserved. # Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional
import numpy as np import numpy as np
import mmcv import mmcv
@ -9,7 +11,10 @@ except ImportError:
torch = None torch = None
def tensor2imgs(tensor, mean=None, std=None, to_rgb=True): def tensor2imgs(tensor,
mean: Optional[tuple] = None,
std: Optional[tuple] = None,
to_rgb: bool = True) -> list:
"""Convert tensor to 3-channel images or 1-channel gray images. """Convert tensor to 3-channel images or 1-channel gray images.
Args: Args: