# 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. import copy import random import warnings from itertools import product from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Union import mmengine import numpy as np import mmcv from mmcv.image.geometric import _scale_size from .base import BaseTransform from .builder import TRANSFORMS from .utils import cache_randomness from .wrappers import Compose Number = Union[int, float] @TRANSFORMS.register_module() class Normalize(BaseTransform): """Normalize the image. Required Keys: - img Modified Keys: - img Added Keys: - img_norm_cfg - mean - std - to_rgb Args: mean (sequence): Mean values of 3 channels. std (sequence): Std values of 3 channels. to_rgb (bool): Whether to convert the image from BGR to RGB before normlizing the image. If ``to_rgb=True``, the order of mean and std should be RGB. If ``to_rgb=False``, the order of mean and std should be the same order of the image. Defaults to True. """ def __init__(self, mean: Sequence[Number], std: Sequence[Number], to_rgb: bool = True) -> None: self.mean = np.array(mean, dtype=np.float32) self.std = np.array(std, dtype=np.float32) self.to_rgb = to_rgb def transform(self, results: dict) -> dict: """Function to normalize images. Args: results (dict): Result dict from loading pipeline. Returns: dict: Normalized results, key 'img_norm_cfg' key is added in to result dict. """ results['img'] = mmcv.imnormalize(results['img'], self.mean, self.std, self.to_rgb) results['img_norm_cfg'] = dict( mean=self.mean, std=self.std, to_rgb=self.to_rgb) return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(mean={self.mean}, std={self.std}, to_rgb={self.to_rgb})' return repr_str @TRANSFORMS.register_module() class Resize(BaseTransform): """Resize images & bbox & seg & keypoints. This transform resizes the input image according to ``scale`` or ``scale_factor``. Bboxes, seg map and keypoints are then resized with the same scale factor. if ``scale`` and ``scale_factor`` are both set, it will use ``scale`` to resize. Required Keys: - img - gt_bboxes (optional) - gt_seg_map (optional) - gt_keypoints (optional) Modified Keys: - img - gt_bboxes - gt_seg_map - gt_keypoints - img_shape Added Keys: - scale - scale_factor - keep_ratio Args: scale (int or tuple): Images scales for resizing. Defaults to None scale_factor (float or tuple[float]): Scale factors for resizing. Defaults to None. keep_ratio (bool): Whether to keep the aspect ratio when resizing the image. Defaults to False. clip_object_border (bool): Whether to clip the objects outside the border of the image. In some dataset like MOT17, the gt bboxes are allowed to cross the border of images. Therefore, we don't need to clip the gt bboxes in these cases. Defaults to True. backend (str): Image resize backend, choices are 'cv2' and 'pillow'. These two backends generates slightly different results. Defaults to 'cv2'. interpolation (str): Interpolation method, accepted values are "nearest", "bilinear", "bicubic", "area", "lanczos" for 'cv2' backend, "nearest", "bilinear" for 'pillow' backend. Defaults to 'bilinear'. """ def __init__(self, scale: Optional[Union[int, Tuple[int, int]]] = None, scale_factor: Optional[Union[float, Tuple[float, float]]] = None, keep_ratio: bool = False, clip_object_border: bool = True, backend: str = 'cv2', interpolation='bilinear') -> None: assert scale is not None or scale_factor is not None, ( '`scale` and' '`scale_factor` can not both be `None`') if scale is None: self.scale = None else: if isinstance(scale, int): self.scale = (scale, scale) else: self.scale = scale self.backend = backend self.interpolation = interpolation self.keep_ratio = keep_ratio self.clip_object_border = clip_object_border if scale_factor is None: self.scale_factor = None elif isinstance(scale_factor, float): self.scale_factor = (scale_factor, scale_factor) elif isinstance(scale_factor, tuple): assert (len(scale_factor)) == 2 self.scale_factor = scale_factor else: raise TypeError( f'expect scale_factor is float or Tuple(float), but' f'get {type(scale_factor)}') def _resize_img(self, results: dict) -> None: """Resize images with ``results['scale']``.""" if results.get('img', None) is not None: if self.keep_ratio: img, scale_factor = mmcv.imrescale( results['img'], results['scale'], interpolation=self.interpolation, return_scale=True, backend=self.backend) # the w_scale and h_scale has minor difference # a real fix should be done in the mmcv.imrescale in the future new_h, new_w = img.shape[:2] h, w = results['img'].shape[:2] w_scale = new_w / w h_scale = new_h / h else: img, w_scale, h_scale = mmcv.imresize( results['img'], results['scale'], interpolation=self.interpolation, return_scale=True, backend=self.backend) results['img'] = img results['img_shape'] = img.shape[:2] results['scale_factor'] = (w_scale, h_scale) results['keep_ratio'] = self.keep_ratio def _resize_bboxes(self, results: dict) -> None: """Resize bounding boxes with ``results['scale_factor']``.""" if results.get('gt_bboxes', None) is not None: bboxes = results['gt_bboxes'] * np.tile( np.array(results['scale_factor']), 2) if self.clip_object_border: bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0, results['img_shape'][1]) bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, results['img_shape'][0]) results['gt_bboxes'] = bboxes def _resize_seg(self, results: dict) -> None: """Resize semantic segmentation map with ``results['scale']``.""" if results.get('gt_seg_map', None) is not None: if self.keep_ratio: gt_seg = mmcv.imrescale( results['gt_seg_map'], results['scale'], interpolation='nearest', backend=self.backend) else: gt_seg = mmcv.imresize( results['gt_seg_map'], results['scale'], interpolation='nearest', backend=self.backend) results['gt_seg_map'] = gt_seg def _resize_keypoints(self, results: dict) -> None: """Resize keypoints with ``results['scale_factor']``.""" if results.get('gt_keypoints', None) is not None: keypoints = results['gt_keypoints'] keypoints[:, :, :2] = keypoints[:, :, :2] * np.array( results['scale_factor']) if self.clip_object_border: keypoints[:, :, 0] = np.clip(keypoints[:, :, 0], 0, results['img_shape'][1]) keypoints[:, :, 1] = np.clip(keypoints[:, :, 1], 0, results['img_shape'][0]) results['gt_keypoints'] = keypoints def transform(self, results: dict) -> dict: """Transform function to resize images, bounding boxes, semantic segmentation map and keypoints. Args: results (dict): Result dict from loading pipeline. Returns: dict: Resized results, 'img', 'gt_bboxes', 'gt_seg_map', 'gt_keypoints', 'scale', 'scale_factor', 'img_shape', and 'keep_ratio' keys are updated in result dict. """ if self.scale: results['scale'] = self.scale else: img_shape = results['img'].shape[:2] results['scale'] = _scale_size(img_shape[::-1], self.scale_factor) # type: ignore self._resize_img(results) self._resize_bboxes(results) self._resize_seg(results) self._resize_keypoints(results) return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(scale={self.scale}, ' repr_str += f'scale_factor={self.scale_factor}, ' repr_str += f'keep_ratio={self.keep_ratio}, ' repr_str += f'clip_object_border={self.clip_object_border}), ' repr_str += f'backend={self.backend}), ' repr_str += f'interpolation={self.interpolation})' return repr_str @TRANSFORMS.register_module() class Pad(BaseTransform): """Pad the image & segmentation map. There are three padding modes: (1) pad to a fixed size and (2) pad to the minimum size that is divisible by some number. and (3)pad to square. Also, pad to square and pad to the minimum size can be used as the same time. Required Keys: - img - gt_bboxes (optional) - gt_seg_map (optional) Modified Keys: - img - gt_seg_map - img_shape Added Keys: - pad_shape - pad_fixed_size - pad_size_divisor Args: size (tuple, optional): Fixed padding size. Expected padding shape (w, h). Defaults to None. size_divisor (int, optional): The divisor of padded size. Defaults to None. pad_to_square (bool): Whether to pad the image into a square. Currently only used for YOLOX. Defaults to False. pad_val (Number | dict[str, Number], optional): Padding value for if the pad_mode is "constant". If it is a single number, the value to pad the image is the number and to pad the semantic segmentation map is 255. If it is a dict, it should have the following keys: - img: The value to pad the image. - seg: The value to pad the semantic segmentation map. Defaults to dict(img=0, seg=255). padding_mode (str): Type of padding. Should be: constant, edge, reflect or symmetric. Defaults to 'constant'. - constant: pads with a constant value, this value is specified with pad_val. - edge: pads with the last value at the edge of the image. - reflect: pads with reflection of image without repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode will result in [3, 2, 1, 2, 3, 4, 3, 2]. - symmetric: pads with reflection of image repeating the last value on the edge. For example, padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode will result in [2, 1, 1, 2, 3, 4, 4, 3] """ def __init__(self, size: Optional[Tuple[int, int]] = None, size_divisor: Optional[int] = None, pad_to_square: bool = False, pad_val: Union[Number, dict] = dict(img=0, seg=255), padding_mode: str = 'constant') -> None: self.size = size self.size_divisor = size_divisor if isinstance(pad_val, int): pad_val = dict(img=pad_val, seg=255) assert isinstance(pad_val, dict), 'pad_val ' self.pad_val = pad_val self.pad_to_square = pad_to_square if pad_to_square: assert size is None, \ 'The size and size_divisor must be None ' \ 'when pad2square is True' else: assert size is not None or size_divisor is not None, \ 'only one of size and size_divisor should be valid' assert size is None or size_divisor is None assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric'] self.padding_mode = padding_mode def _pad_img(self, results: dict) -> None: """Pad images according to ``self.size``.""" pad_val = self.pad_val.get('img', 0) size = None if self.pad_to_square: max_size = max(results['img'].shape[:2]) size = (max_size, max_size) if self.size_divisor is not None: if size is None: size = (results['img'].shape[0], results['img'].shape[1]) pad_h = int(np.ceil( size[0] / self.size_divisor)) * self.size_divisor pad_w = int(np.ceil( size[1] / self.size_divisor)) * self.size_divisor size = (pad_h, pad_w) elif self.size is not None: size = self.size[::-1] if isinstance(pad_val, int) and results['img'].ndim == 3: pad_val = tuple(pad_val for _ in range(results['img'].shape[2])) padded_img = mmcv.impad( results['img'], shape=size, pad_val=pad_val, padding_mode=self.padding_mode) results['img'] = padded_img results['pad_shape'] = padded_img.shape results['pad_fixed_size'] = self.size results['pad_size_divisor'] = self.size_divisor results['img_shape'] = padded_img.shape[:2] def _pad_seg(self, results: dict) -> None: """Pad semantic segmentation map according to ``results['pad_shape']``.""" if results.get('gt_seg_map', None) is not None: pad_val = self.pad_val.get('seg', 255) if isinstance(pad_val, int) and results['gt_seg_map'].ndim == 3: pad_val = tuple( pad_val for _ in range(results['gt_seg_map'].shape[2])) results['gt_seg_map'] = mmcv.impad( results['gt_seg_map'], shape=results['pad_shape'][:2], pad_val=pad_val, padding_mode=self.padding_mode) def transform(self, results: dict) -> dict: """Call function to pad images, masks, semantic segmentation maps. Args: results (dict): Result dict from loading pipeline. Returns: dict: Updated result dict. """ self._pad_img(results) self._pad_seg(results) return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(size={self.size}, ' repr_str += f'size_divisor={self.size_divisor}, ' repr_str += f'pad_to_square={self.pad_to_square}, ' repr_str += f'pad_val={self.pad_val}), ' repr_str += f'padding_mode={self.padding_mode})' return repr_str @TRANSFORMS.register_module() class CenterCrop(BaseTransform): """Crop the center of the image, segmentation masks, bounding boxes and key points. If the crop area exceeds the original image and ``auto_pad`` is True, the original image will be padded before cropping. Required Keys: - img - gt_seg_map (optional) - gt_bboxes (optional) - gt_keypoints (optional) Modified Keys: - img - img_shape - gt_seg_map (optional) - gt_bboxes (optional) - gt_keypoints (optional) Added Key: - pad_shape Args: crop_size (Union[int, Tuple[int, int]]): Expected size after cropping with the format of (w, h). If set to an integer, then cropping width and height are equal to this integer. auto_pad (bool): Whether to pad the image if it's smaller than the ``crop_size``. Defaults to False. pad_cfg (dict): Base config for padding. Refer to ``mmcv.Pad`` for detail. Defaults to ``dict(type='Pad')``. clip_object_border (bool): Whether to clip the objects outside the border of the image. In some dataset like MOT17, the gt bboxes are allowed to cross the border of images. Therefore, we don't need to clip the gt bboxes in these cases. Defaults to True. """ def __init__(self, crop_size: Union[int, Tuple[int, int]], auto_pad: bool = False, pad_cfg: dict = dict(type='Pad'), clip_object_border: bool = True) -> None: super().__init__() assert isinstance(crop_size, int) or ( isinstance(crop_size, tuple) and len(crop_size) == 2 ), 'The expected crop_size is an integer, or a tuple containing two ' 'intergers' if isinstance(crop_size, int): crop_size = (crop_size, crop_size) assert crop_size[0] > 0 and crop_size[1] > 0 self.crop_size = crop_size self.auto_pad = auto_pad self.pad_cfg = pad_cfg.copy() # size will be overwritten if 'size' in self.pad_cfg and auto_pad: warnings.warn('``size`` is set in ``pad_cfg``,' 'however this argument will be overwritten' ' according to crop size and image size') self.clip_object_border = clip_object_border def _crop_img(self, results: dict, bboxes: np.ndarray) -> None: """Crop image. Args: results (dict): Result dict contains the data to transform. bboxes (np.ndarray): Shape (4, ), location of cropped bboxes. """ if results.get('img', None) is not None: img = mmcv.imcrop(results['img'], bboxes=bboxes) img_shape = img.shape[:2] # type: ignore results['img'] = img results['img_shape'] = img_shape results['pad_shape'] = img_shape def _crop_seg_map(self, results: dict, bboxes: np.ndarray) -> None: """Crop semantic segmentation map. Args: results (dict): Result dict contains the data to transform. bboxes (np.ndarray): Shape (4, ), location of cropped bboxes. """ if results.get('gt_seg_map', None) is not None: img = mmcv.imcrop(results['gt_seg_map'], bboxes=bboxes) results['gt_seg_map'] = img def _crop_bboxes(self, results: dict, bboxes: np.ndarray) -> None: """Update bounding boxes according to CenterCrop. Args: results (dict): Result dict contains the data to transform. bboxes (np.ndarray): Shape (4, ), location of cropped bboxes. """ if 'gt_bboxes' in results: offset_w = bboxes[0] offset_h = bboxes[1] bbox_offset = np.array([offset_w, offset_h, offset_w, offset_h]) # gt_bboxes has shape (num_gts, 4) in (tl_x, tl_y, br_x, br_y) # order. gt_bboxes = results['gt_bboxes'] - bbox_offset if self.clip_object_border: gt_bboxes[:, 0::2] = np.clip(gt_bboxes[:, 0::2], 0, results['img'].shape[1]) gt_bboxes[:, 1::2] = np.clip(gt_bboxes[:, 1::2], 0, results['img'].shape[0]) results['gt_bboxes'] = gt_bboxes def _crop_keypoints(self, results: dict, bboxes: np.ndarray) -> None: """Update key points according to CenterCrop. Keypoints that not in the cropped image will be set invisible. Args: results (dict): Result dict contains the data to transform. bboxes (np.ndarray): Shape (4, ), location of cropped bboxes. """ if 'gt_keypoints' in results: offset_w = bboxes[0] offset_h = bboxes[1] keypoints_offset = np.array([offset_w, offset_h, 0]) # gt_keypoints has shape (N, NK, 3) in (x, y, visibility) order, # NK = number of points per object gt_keypoints = results['gt_keypoints'] - keypoints_offset # set gt_kepoints out of the result image invisible height, width = results['img'].shape[:2] valid_pos = (gt_keypoints[:, :, 0] >= 0) * (gt_keypoints[:, :, 0] < width) * (gt_keypoints[:, :, 1] >= 0) * ( gt_keypoints[:, :, 1] < height) gt_keypoints[:, :, 2] = np.where(valid_pos, gt_keypoints[:, :, 2], 0) gt_keypoints[:, :, 0] = np.clip(gt_keypoints[:, :, 0], 0, results['img'].shape[1]) gt_keypoints[:, :, 1] = np.clip(gt_keypoints[:, :, 1], 0, results['img'].shape[0]) results['gt_keypoints'] = gt_keypoints def transform(self, results: dict) -> dict: """Apply center crop on results. Args: results (dict): Result dict contains the data to transform. Returns: dict: Results with CenterCropped image and semantic segmentation map. """ crop_width, crop_height = self.crop_size[0], self.crop_size[1] assert 'img' in results, '`img` is not found in results' img = results['img'] # img.shape has length 2 for grayscale, length 3 for color img_height, img_width = img.shape[:2] if crop_height > img_height or crop_width > img_width: if self.auto_pad: # pad the area img_height = max(img_height, crop_height) img_width = max(img_width, crop_width) pad_size = (img_width, img_height) _pad_cfg = self.pad_cfg.copy() _pad_cfg.update(dict(size=pad_size)) pad_transform = TRANSFORMS.build(_pad_cfg) results = pad_transform(results) else: crop_height = min(crop_height, img_height) crop_width = min(crop_width, img_width) y1 = max(0, int(round((img_height - crop_height) / 2.))) x1 = max(0, int(round((img_width - crop_width) / 2.))) y2 = min(img_height, y1 + crop_height) - 1 x2 = min(img_width, x1 + crop_width) - 1 bboxes = np.array([x1, y1, x2, y2]) # crop the image self._crop_img(results, bboxes) # crop the gt_seg_map self._crop_seg_map(results, bboxes) # crop the bounding box self._crop_bboxes(results, bboxes) # crop the keypoints self._crop_keypoints(results, bboxes) return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(crop_size = {self.crop_size}' repr_str += f', auto_pad={self.auto_pad}' repr_str += f', pad_cfg={self.pad_cfg}' repr_str += f',clip_object_border = {self.clip_object_border})' return repr_str @TRANSFORMS.register_module() class RandomGrayscale(BaseTransform): """Randomly convert image to grayscale with a probability. Required Key: - img Modified Key: - img Added Keys: - grayscale - grayscale_weights Args: prob (float): Probability that image should be converted to grayscale. Defaults to 0.1. keep_channels (bool): Whether keep channel number the same as input. Defaults to False. channel_weights (tuple): The grayscale weights of each channel, and the weights will be normalized. For example, (1, 2, 1) will be normalized as (0.25, 0.5, 0.25). Defaults to (1., 1., 1.). color_format (str): Color format set to be any of 'bgr', 'rgb', 'hsv'. Note: 'hsv' image will be transformed into 'bgr' format no matter whether it is grayscaled. Defaults to 'bgr'. """ def __init__(self, prob: float = 0.1, keep_channels: bool = False, channel_weights: Sequence[float] = (1., 1., 1.), color_format: str = 'bgr') -> None: super().__init__() assert 0. <= prob <= 1., ('The range of ``prob`` value is [0., 1.],' + f' but got {prob} instead') self.prob = prob self.keep_channels = keep_channels self.channel_weights = channel_weights assert color_format in ['bgr', 'rgb', 'hsv'] self.color_format = color_format @cache_randomness def _random_prob(self): return random.random() def transform(self, results: dict) -> dict: """Apply random grayscale on results. Args: results (dict): Result dict contains the data to transform. Returns: dict: Results with grayscale image. """ img = results['img'] # convert hsv to bgr if self.color_format == 'hsv': img = mmcv.hsv2bgr(img) img = img[..., None] if img.ndim == 2 else img num_output_channels = img.shape[2] if self._random_prob() < self.prob: if num_output_channels > 1: assert num_output_channels == len( self.channel_weights ), 'The length of ``channel_weights`` are supposed to be ' f'num_output_channels, but got {len(self.channel_weights)}' ' instead.' normalized_weights = ( np.array(self.channel_weights) / sum(self.channel_weights)) img = (normalized_weights * img).sum(axis=2) img = img.astype('uint8') if self.keep_channels: img = img[:, :, None] results['img'] = np.dstack( [img for _ in range(num_output_channels)]) else: results['img'] = img return results img = img.astype('uint8') results['img'] = img return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(prob = {self.prob}' repr_str += f', keep_channels = {self.keep_channels}' repr_str += f', channel_weights = {self.channel_weights}' repr_str += f', color_format = {self.color_format})' return repr_str @TRANSFORMS.register_module() class MultiScaleFlipAug(BaseTransform): """Test-time augmentation with multiple scales and flipping. An example configuration is as followed: .. code-block:: dict( type='MultiScaleFlipAug', scales=[(1333, 400), (1333, 800)], flip=True, transforms=[ dict(type='Normalize', **img_norm_cfg), dict(type='Pad', size_divisor=1), dict(type='ImageToTensor', keys=['img']), dict(type='Collect', keys=['img']) ]) ``results`` will be resized using all the sizes in ``scales``. If ``flip`` is True, then flipped results will also be added into output list. For the above configuration, there are four combinations of resize and flip: - Resize to (1333, 400) + no flip - Resize to (1333, 400) + flip - Resize to (1333, 800) + no flip - resize to (1333, 800) + flip The four results are then transformed with ``transforms`` argument. After that, results are wrapped into lists of the same length as below: .. code-block:: dict( inputs=[...], data_samples=[...] ) Where the length of ``inputs`` and ``data_samples`` are both 4. Required Keys: - Depending on the requirements of the ``transforms`` parameter. Modified Keys: - All output keys of each transform. Args: transforms (list[dict]): Transforms to be applied to each resized and flipped data. scales (tuple | list[tuple] | None): Images scales for resizing. scale_factor (float or tuple[float]): Scale factors for resizing. Defaults to None. allow_flip (bool): Whether apply flip augmentation. Defaults to False. flip_direction (str | list[str]): Flip augmentation directions, options are "horizontal", "vertical" and "diagonal". If flip_direction is a list, multiple flip augmentations will be applied. It has no effect when flip == False. Defaults to "horizontal". resize_cfg (dict): Base config for resizing. Defaults to ``dict(type='Resize', keep_ratio=True)``. flip_cfg (dict): Base config for flipping. Defaults to ``dict(type='RandomFlip')``. """ def __init__( self, transforms: List[dict], scales: Optional[Union[Tuple, List[Tuple]]] = None, scale_factor: Optional[Union[float, List[float]]] = None, allow_flip: bool = False, flip_direction: Union[str, List[str]] = 'horizontal', resize_cfg: dict = dict(type='Resize', keep_ratio=True), flip_cfg: dict = dict(type='RandomFlip') ) -> None: super().__init__() self.transforms = Compose(transforms) # type: ignore if scales is not None: self.scales = scales if isinstance(scales, list) else [scales] self.scale_key = 'scale' assert mmengine.is_list_of(self.scales, tuple) else: # if ``scales`` and ``scale_factor`` both be ``None`` if scale_factor is None: self.scales = [1.] # type: ignore elif isinstance(scale_factor, list): self.scales = scale_factor # type: ignore else: self.scales = [scale_factor] # type: ignore self.scale_key = 'scale_factor' self.allow_flip = allow_flip self.flip_direction = flip_direction if isinstance( flip_direction, list) else [flip_direction] assert mmengine.is_list_of(self.flip_direction, str) if not self.allow_flip and self.flip_direction != ['horizontal']: warnings.warn( 'flip_direction has no effect when flip is set to False') self.resize_cfg = resize_cfg.copy() self.flip_cfg = flip_cfg def transform(self, results: dict) -> Dict: """Apply test time augment transforms on results. Args: results (dict): Result dict contains the data to transform. Returns: dict: The augmented data, where each value is wrapped into a list. """ data_samples = [] inputs = [] flip_args = [(False, '')] if self.allow_flip: flip_args += [(True, direction) for direction in self.flip_direction] for scale in self.scales: for flip, direction in flip_args: _resize_cfg = self.resize_cfg.copy() _resize_cfg.update({self.scale_key: scale}) _resize_flip = [_resize_cfg] if flip: _flip_cfg = self.flip_cfg.copy() _flip_cfg.update(prob=1.0, direction=direction) _resize_flip.append(_flip_cfg) else: results['flip'] = False results['flip_direction'] = None resize_flip = Compose(_resize_flip) _results = resize_flip(results.copy()) packed_results = self.transforms(_results) # type: ignore inputs.append(packed_results['inputs']) # type: ignore data_samples.append( packed_results['data_sample']) # type: ignore return dict(inputs=inputs, data_sample=data_samples) def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(transforms={self.transforms}' repr_str += f', scales={self.scales}' repr_str += f', allow_flip={self.allow_flip}' repr_str += f', flip_direction={self.flip_direction})' return repr_str @TRANSFORMS.register_module() class TestTimeAug(BaseTransform): """Test-time augmentation transform. An example configuration is as followed: .. code-block:: dict(type='TestTimeAug', transforms=[ [dict(type='Resize', scale=(1333, 400), keep_ratio=True), dict(type='Resize', scale=(1333, 800), keep_ratio=True)], [dict(type='RandomFlip', prob=1.), dict(type='RandomFlip', prob=0.)], [dict(type='PackDetInputs', meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', 'scale_factor', 'flip', 'flip_direction'))]]) ``results`` will be transformed using all transforms defined in ``transforms`` arguments. For the above configuration, there are four combinations of resize and flip: - Resize to (1333, 400) + no flip - Resize to (1333, 400) + flip - Resize to (1333, 800) + no flip - resize to (1333, 800) + flip After that, results are wrapped into lists of the same length as below: .. code-block:: dict( inputs=[...], data_samples=[...] ) The length of ``inputs`` and ``data_samples`` are both 4. Required Keys: - Depending on the requirements of the ``transforms`` parameter. Modified Keys: - All output keys of each transform. Args: transforms (list[list[dict]]): Transforms to be applied to data sampled from dataset. ``transforms`` is a list of list, and each list element usually represents a series of transforms with the same type and different arguments. Data will be processed by each list elements sequentially. See more information in :meth:`transform`. """ def __init__(self, transforms: list): for i, transform_list in enumerate(transforms): for j, transform in enumerate(transform_list): if isinstance(transform, dict): transform_list[j] = TRANSFORMS.build(transform) elif callable(transform): continue else: raise TypeError( 'transform must be callable or a dict, but got' f' {type(transform)}') transforms[i] = transform_list self.subroutines = [ Compose(subroutine) for subroutine in product(*transforms) ] def transform(self, results: dict) -> dict: """Apply all transforms defined in :attr:`transforms` to the results. As the example given in :obj:`TestTimeAug`, ``transforms`` consists of 2 ``Resize``, 2 ``RandomFlip`` and 1 ``PackDetInputs``. The data sampled from dataset will be processed as follows: 1. Data will be processed by 2 ``Resize`` and return a list of 2 results. 2. Each result in list will be further passed to 2 ``RandomFlip``, and aggregates into a list of 4 results. 3. Each result will be processed by ``PackDetInputs``, and return a list of dict. 4. Aggregates the same fields of results, and finally returns a dict. Each value of the dict represents 4 transformed results. Args: results (dict): Result dict contains the data to transform. Returns: dict: The augmented data, where each value is wrapped into a list. """ results_list = [] # type: ignore for subroutine in self.subroutines: result = subroutine(copy.deepcopy(results)) assert isinstance(result, dict), ( f'Data processed by {subroutine} must return a dict, but got ' f'{result}') assert result is not None, ( f'Data processed by {subroutine} in `TestTimeAug` should not ' 'be None! Please check your validation dataset and the ' f'transforms in {subroutine}') results_list.append(result) aug_data_dict = { key: [item[key] for item in results_list] # type: ignore for key in results_list[0] # type: ignore } return aug_data_dict def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += 'transforms=\n' for subroutine in self.subroutines: repr_str += f'{repr(subroutine)}\n' return repr_str @TRANSFORMS.register_module() class RandomChoiceResize(BaseTransform): """Resize images & bbox & mask from a list of multiple scales. This transform resizes the input image to some scale. Bboxes and masks are then resized with the same scale factor. Resize scale will be randomly selected from ``scales``. How to choose the target scale to resize the image will follow the rules below: - if `scale` is a list of tuple, the target scale is sampled from the list uniformally. - if `scale` is a tuple, the target scale will be set to the tuple. Required Keys: - img - gt_bboxes (optional) - gt_seg_map (optional) - gt_keypoints (optional) Modified Keys: - img - img_shape - gt_bboxes (optional) - gt_seg_map (optional) - gt_keypoints (optional) Added Keys: - scale - scale_factor - scale_idx - keep_ratio Args: scales (Union[list, Tuple]): Images scales for resizing. resize_type (str): The type of resize class to use. Defaults to "Resize". **resize_kwargs: Other keyword arguments for the ``resize_type``. Note: By defaults, the ``resize_type`` is "Resize", if it's not overwritten by your registry, it indicates the :class:`mmcv.Resize`. And therefore, ``resize_kwargs`` accepts any keyword arguments of it, like ``keep_ratio``, ``interpolation`` and so on. If you want to use your custom resize class, the class should accept ``scale`` argument and have ``scale`` attribution which determines the resize shape. """ def __init__( self, scales: Sequence[Union[int, Tuple]], resize_type: str = 'Resize', **resize_kwargs, ) -> None: super().__init__() if isinstance(scales, list): self.scales = scales else: self.scales = [scales] assert mmengine.is_seq_of(self.scales, (tuple, int)) self.resize_cfg = dict(type=resize_type, **resize_kwargs) # create a empty Resize object self.resize = TRANSFORMS.build({'scale': 0, **self.resize_cfg}) @cache_randomness def _random_select(self) -> Tuple[int, int]: """Randomly select an scale from given candidates. Returns: (tuple, int): Returns a tuple ``(scale, scale_dix)``, where ``scale`` is the selected image scale and ``scale_idx`` is the selected index in the given candidates. """ scale_idx = np.random.randint(len(self.scales)) scale = self.scales[scale_idx] return scale, scale_idx def transform(self, results: dict) -> dict: """Apply resize transforms on results from a list of scales. Args: results (dict): Result dict contains the data to transform. Returns: dict: Resized results, 'img', 'gt_bboxes', 'gt_seg_map', 'gt_keypoints', 'scale', 'scale_factor', 'img_shape', and 'keep_ratio' keys are updated in result dict. """ target_scale, scale_idx = self._random_select() self.resize.scale = target_scale results = self.resize(results) results['scale_idx'] = scale_idx return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(scales={self.scales}' repr_str += f', resize_cfg={self.resize_cfg})' return repr_str @TRANSFORMS.register_module() class RandomFlip(BaseTransform): """Flip the image & bbox & keypoints & segmentation map. Added or Updated keys: flip, flip_direction, img, gt_bboxes, gt_seg_map, and gt_keypoints. There are 3 flip modes: - ``prob`` is float, ``direction`` is string: the image will be ``direction``ly flipped with probability of ``prob`` . E.g., ``prob=0.5``, ``direction='horizontal'``, then image will be horizontally flipped with probability of 0.5. - ``prob`` is float, ``direction`` is list of string: the image will be ``direction[i]``ly flipped with probability of ``prob/len(direction)``. E.g., ``prob=0.5``, ``direction=['horizontal', 'vertical']``, then image will be horizontally flipped with probability of 0.25, vertically with probability of 0.25. - ``prob`` is list of float, ``direction`` is list of string: given ``len(prob) == len(direction)``, the image will be ``direction[i]``ly flipped with probability of ``prob[i]``. E.g., ``prob=[0.3, 0.5]``, ``direction=['horizontal', 'vertical']``, then image will be horizontally flipped with probability of 0.3, vertically with probability of 0.5. Required Keys: - img - gt_bboxes (optional) - gt_seg_map (optional) - gt_keypoints (optional) Modified Keys: - img - gt_bboxes (optional) - gt_seg_map (optional) - gt_keypoints (optional) Added Keys: - flip - flip_direction - swap_seg_labels (optional) Args: prob (float | list[float], optional): The flipping probability. Defaults to None. direction(str | list[str]): The flipping direction. Options If input is a list, the length must equal ``prob``. Each element in ``prob`` indicates the flip probability of corresponding direction. Defaults to 'horizontal'. swap_seg_labels (list, optional): The label pair need to be swapped for ground truth, like 'left arm' and 'right arm' need to be swapped after horizontal flipping. For example, ``[(1, 5)]``, where 1/5 is the label of the left/right arm. Defaults to None. """ def __init__(self, prob: Optional[Union[float, Iterable[float]]] = None, direction: Union[str, Sequence[Optional[str]]] = 'horizontal', swap_seg_labels: Optional[Sequence] = None) -> None: if isinstance(prob, list): assert mmengine.is_list_of(prob, float) assert 0 <= sum(prob) <= 1 elif isinstance(prob, float): assert 0 <= prob <= 1 else: raise ValueError(f'probs must be float or list of float, but \ got `{type(prob)}`.') self.prob = prob self.swap_seg_labels = swap_seg_labels valid_directions = ['horizontal', 'vertical', 'diagonal'] if isinstance(direction, str): assert direction in valid_directions elif isinstance(direction, list): assert mmengine.is_list_of(direction, str) assert set(direction).issubset(set(valid_directions)) else: raise ValueError(f'direction must be either str or list of str, \ but got `{type(direction)}`.') self.direction = direction if isinstance(prob, list): assert len(prob) == len(self.direction) def _flip_bbox(self, bboxes: np.ndarray, img_shape: Tuple[int, int], direction: str) -> np.ndarray: """Flip bboxes horizontally. Args: bboxes (numpy.ndarray): Bounding boxes, shape (..., 4*k) img_shape (tuple[int]): Image shape (height, width) direction (str): Flip direction. Options are 'horizontal', 'vertical', and 'diagonal'. Returns: numpy.ndarray: Flipped bounding boxes. """ assert bboxes.shape[-1] % 4 == 0 flipped = bboxes.copy() h, w = img_shape if direction == 'horizontal': flipped[..., 0::4] = w - bboxes[..., 2::4] flipped[..., 2::4] = w - bboxes[..., 0::4] elif direction == 'vertical': flipped[..., 1::4] = h - bboxes[..., 3::4] flipped[..., 3::4] = h - bboxes[..., 1::4] elif direction == 'diagonal': flipped[..., 0::4] = w - bboxes[..., 2::4] flipped[..., 1::4] = h - bboxes[..., 3::4] flipped[..., 2::4] = w - bboxes[..., 0::4] flipped[..., 3::4] = h - bboxes[..., 1::4] else: raise ValueError( f"Flipping direction must be 'horizontal', 'vertical', \ or 'diagonal', but got '{direction}'") return flipped def _flip_keypoints( self, keypoints: np.ndarray, img_shape: Tuple[int, int], direction: str, ) -> np.ndarray: """Flip keypoints horizontally, vertically or diagonally. Args: keypoints (numpy.ndarray): Keypoints, shape (..., 2) img_shape (tuple[int]): Image shape (height, width) direction (str): Flip direction. Options are 'horizontal', 'vertical', and 'diagonal'. Returns: numpy.ndarray: Flipped keypoints. """ meta_info = keypoints[..., 2:] keypoints = keypoints[..., :2] flipped = keypoints.copy() h, w = img_shape if direction == 'horizontal': flipped[..., 0::2] = w - keypoints[..., 0::2] elif direction == 'vertical': flipped[..., 1::2] = h - keypoints[..., 1::2] elif direction == 'diagonal': flipped[..., 0::2] = w - keypoints[..., 0::2] flipped[..., 1::2] = h - keypoints[..., 1::2] else: raise ValueError( f"Flipping direction must be 'horizontal', 'vertical', \ or 'diagonal', but got '{direction}'") flipped = np.concatenate([flipped, meta_info], axis=-1) return flipped def _flip_seg_map(self, seg_map: dict, direction: str) -> np.ndarray: """Flip segmentation map horizontally, vertically or diagonally. Args: seg_map (numpy.ndarray): segmentation map, shape (H, W). direction (str): Flip direction. Options are 'horizontal', 'vertical'. Returns: numpy.ndarray: Flipped segmentation map. """ seg_map = mmcv.imflip(seg_map, direction=direction) if self.swap_seg_labels is not None: # to handle datasets with left/right annotations # like 'Left-arm' and 'Right-arm' in LIP dataset # Modified from https://github.com/openseg-group/openseg.pytorch/blob/master/lib/datasets/tools/cv2_aug_transforms.py # noqa:E501 # Licensed under MIT license temp = seg_map.copy() assert isinstance(self.swap_seg_labels, (tuple, list)) for pair in self.swap_seg_labels: assert isinstance(pair, (tuple, list)) and len(pair) == 2, \ 'swap_seg_labels must be a sequence with pair, but got ' \ f'{self.swap_seg_labels}.' seg_map[temp == pair[0]] = pair[1] seg_map[temp == pair[1]] = pair[0] return seg_map @cache_randomness def _choose_direction(self) -> str: """Choose the flip direction according to `prob` and `direction`""" if isinstance(self.direction, Sequence) and not isinstance(self.direction, str): # None means non-flip direction_list: list = list(self.direction) + [None] elif isinstance(self.direction, str): # None means non-flip direction_list = [self.direction, None] if isinstance(self.prob, list): non_prob: float = 1 - sum(self.prob) prob_list = self.prob + [non_prob] elif isinstance(self.prob, float): non_prob = 1. - self.prob # exclude non-flip single_ratio = self.prob / (len(direction_list) - 1) prob_list = [single_ratio] * (len(direction_list) - 1) + [non_prob] cur_dir = np.random.choice(direction_list, p=prob_list) return cur_dir def _flip(self, results: dict) -> None: """Flip images, bounding boxes, semantic segmentation map and keypoints.""" # flip image results['img'] = mmcv.imflip( results['img'], direction=results['flip_direction']) img_shape = results['img'].shape[:2] # flip bboxes if results.get('gt_bboxes', None) is not None: results['gt_bboxes'] = self._flip_bbox(results['gt_bboxes'], img_shape, results['flip_direction']) # flip keypoints if results.get('gt_keypoints', None) is not None: results['gt_keypoints'] = self._flip_keypoints( results['gt_keypoints'], img_shape, results['flip_direction']) # flip seg map if results.get('gt_seg_map', None) is not None: results['gt_seg_map'] = self._flip_seg_map( results['gt_seg_map'], direction=results['flip_direction']) results['swap_seg_labels'] = self.swap_seg_labels def _flip_on_direction(self, results: dict) -> None: """Function to flip images, bounding boxes, semantic segmentation map and keypoints.""" cur_dir = self._choose_direction() if cur_dir is None: results['flip'] = False results['flip_direction'] = None else: results['flip'] = True results['flip_direction'] = cur_dir self._flip(results) def transform(self, results: dict) -> dict: """Transform function to flip images, bounding boxes, semantic segmentation map and keypoints. Args: results (dict): Result dict from loading pipeline. Returns: dict: Flipped results, 'img', 'gt_bboxes', 'gt_seg_map', 'gt_keypoints', 'flip', and 'flip_direction' keys are updated in result dict. """ self._flip_on_direction(results) return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(prob={self.prob}, ' repr_str += f'direction={self.direction})' return repr_str @TRANSFORMS.register_module() class RandomResize(BaseTransform): """Random resize images & bbox & keypoints. How to choose the target scale to resize the image will follow the rules below: - if ``scale`` is a sequence of tuple .. math:: target\\_scale[0] \\sim Uniform([scale[0][0], scale[1][0]]) .. math:: target\\_scale[1] \\sim Uniform([scale[0][1], scale[1][1]]) Following the resize order of weight and height in cv2, ``scale[i][0]`` is for width, and ``scale[i][1]`` is for height. - if ``scale`` is a tuple .. math:: target\\_scale[0] \\sim Uniform([ratio\\_range[0], ratio\\_range[1]]) * scale[0] .. math:: target\\_scale[1] \\sim Uniform([ratio\\_range[0], ratio\\_range[1]]) * scale[1] Following the resize order of weight and height in cv2, ``ratio_range[0]`` is for width, and ``ratio_range[1]`` is for height. - if ``keep_ratio`` is True, the minimum value of ``target_scale`` will be used to set the shorter side and the maximum value will be used to set the longer side. - if ``keep_ratio`` is False, the value of ``target_scale`` will be used to reisze the width and height accordingly. Required Keys: - img - gt_bboxes - gt_seg_map - gt_keypoints Modified Keys: - img - gt_bboxes - gt_seg_map - gt_keypoints - img_shape Added Keys: - scale - scale_factor - keep_ratio Args: scale (tuple or Sequence[tuple]): Images scales for resizing. Defaults to None. ratio_range (tuple[float], optional): (min_ratio, max_ratio). Defaults to None. resize_type (str): The type of resize class to use. Defaults to "Resize". **resize_kwargs: Other keyword arguments for the ``resize_type``. Note: By defaults, the ``resize_type`` is "Resize", if it's not overwritten by your registry, it indicates the :class:`mmcv.Resize`. And therefore, ``resize_kwargs`` accepts any keyword arguments of it, like ``keep_ratio``, ``interpolation`` and so on. If you want to use your custom resize class, the class should accept ``scale`` argument and have ``scale`` attribution which determines the resize shape. """ def __init__( self, scale: Union[Tuple[int, int], Sequence[Tuple[int, int]]], ratio_range: Tuple[float, float] = None, resize_type: str = 'Resize', **resize_kwargs, ) -> None: self.scale = scale self.ratio_range = ratio_range self.resize_cfg = dict(type=resize_type, **resize_kwargs) # create a empty Reisize object self.resize = TRANSFORMS.build({'scale': 0, **self.resize_cfg}) @staticmethod def _random_sample(scales: Sequence[Tuple[int, int]]) -> tuple: """Private function to randomly sample a scale from a list of tuples. Args: scales (list[tuple]): Images scale range for sampling. There must be two tuples in scales, which specify the lower and upper bound of image scales. Returns: tuple: The targeted scale of the image to be resized. """ assert mmengine.is_list_of(scales, tuple) and len(scales) == 2 scale_0 = [scales[0][0], scales[1][0]] scale_1 = [scales[0][1], scales[1][1]] edge_0 = np.random.randint(min(scale_0), max(scale_0) + 1) edge_1 = np.random.randint(min(scale_1), max(scale_1) + 1) scale = (edge_0, edge_1) return scale @staticmethod def _random_sample_ratio(scale: tuple, ratio_range: Tuple[float, float]) -> tuple: """Private function to randomly sample a scale from a tuple. A ratio will be randomly sampled from the range specified by ``ratio_range``. Then it would be multiplied with ``scale`` to generate sampled scale. Args: scale (tuple): Images scale base to multiply with ratio. ratio_range (tuple[float]): The minimum and maximum ratio to scale the ``scale``. Returns: tuple: The targeted scale of the image to be resized. """ assert isinstance(scale, tuple) and len(scale) == 2 min_ratio, max_ratio = ratio_range assert min_ratio <= max_ratio ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio scale = int(scale[0] * ratio), int(scale[1] * ratio) return scale @cache_randomness def _random_scale(self) -> tuple: """Private function to randomly sample an scale according to the type of ``scale``. Returns: tuple: The targeted scale of the image to be resized. """ if mmengine.is_tuple_of(self.scale, int): assert self.ratio_range is not None and len(self.ratio_range) == 2 scale = self._random_sample_ratio( self.scale, # type: ignore self.ratio_range) elif mmengine.is_seq_of(self.scale, tuple): scale = self._random_sample(self.scale) # type: ignore else: raise NotImplementedError('Do not support sampling function ' f'for "{self.scale}"') return scale def transform(self, results: dict) -> dict: """Transform function to resize images, bounding boxes, semantic segmentation map. Args: results (dict): Result dict from loading pipeline. Returns: dict: Resized results, ``img``, ``gt_bboxes``, ``gt_semantic_seg``, ``gt_keypoints``, ``scale``, ``scale_factor``, ``img_shape``, and ``keep_ratio`` keys are updated in result dict. """ results['scale'] = self._random_scale() self.resize.scale = results['scale'] results = self.resize(results) return results def __repr__(self) -> str: repr_str = self.__class__.__name__ repr_str += f'(scale={self.scale}, ' repr_str += f'ratio_range={self.ratio_range}, ' repr_str += f'resize_cfg={self.resize_cfg})' return repr_str