Spaces:
Runtime error
Runtime error
File size: 3,683 Bytes
2366e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
# Copyright (c) OpenMMLab. All rights reserved.
import mmcv
import numpy as np
from mmdet.datasets.builder import PIPELINES
from mmdet.datasets.pipelines.compose import Compose
@PIPELINES.register_module()
class MultiRotateAugOCR:
"""Test-time augmentation with multiple rotations in the case that
img_height > img_width.
An example configuration is as follows:
.. code-block::
rotate_degrees=[0, 90, 270],
transforms=[
dict(
type='ResizeOCR',
height=32,
min_width=32,
max_width=160,
keep_aspect_ratio=True),
dict(type='ToTensorOCR'),
dict(type='NormalizeOCR', **img_norm_cfg),
dict(
type='Collect',
keys=['img'],
meta_keys=[
'filename', 'ori_shape', 'img_shape', 'valid_ratio'
]),
]
After MultiRotateAugOCR with above configuration, the results are wrapped
into lists of the same length as follows:
.. code-block::
dict(
img=[...],
img_shape=[...]
...
)
Args:
transforms (list[dict]): Transformation applied for each augmentation.
rotate_degrees (list[int] | None): Degrees of anti-clockwise rotation.
force_rotate (bool): If True, rotate image by 'rotate_degrees'
while ignore image aspect ratio.
"""
def __init__(self, transforms, rotate_degrees=None, force_rotate=False):
self.transforms = Compose(transforms)
self.force_rotate = force_rotate
if rotate_degrees is not None:
self.rotate_degrees = rotate_degrees if isinstance(
rotate_degrees, list) else [rotate_degrees]
assert mmcv.is_list_of(self.rotate_degrees, int)
for degree in self.rotate_degrees:
assert 0 <= degree < 360
assert degree % 90 == 0
if 0 not in self.rotate_degrees:
self.rotate_degrees.append(0)
else:
self.rotate_degrees = [0]
def __call__(self, results):
"""Call function to apply test time augment transformation to results.
Args:
results (dict): Result dict contains the data to be transformed.
Returns:
dict[str: list]: The augmented data, where each value is wrapped
into a list.
"""
img_shape = results['img_shape']
ori_height, ori_width = img_shape[:2]
if not self.force_rotate and ori_height <= ori_width:
rotate_degrees = [0]
else:
rotate_degrees = self.rotate_degrees
aug_data = []
for degree in set(rotate_degrees):
_results = results.copy()
if degree == 0:
pass
elif degree == 90:
_results['img'] = np.rot90(_results['img'], 1)
elif degree == 180:
_results['img'] = np.rot90(_results['img'], 2)
elif degree == 270:
_results['img'] = np.rot90(_results['img'], 3)
data = self.transforms(_results)
aug_data.append(data)
# list of dict to dict of list
aug_data_dict = {key: [] for key in aug_data[0]}
for data in aug_data:
for key, val in data.items():
aug_data_dict[key].append(val)
return aug_data_dict
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'(transforms={self.transforms}, '
repr_str += f'rotate_degrees={self.rotate_degrees})'
return repr_str
|