Spaces:
Running
Running
File size: 9,345 Bytes
5de2f8f |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This code is refer from:
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/iaa_augment.py
"""
import os
# Prevent automatic updates in Albumentations for stability in augmentation behavior
os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1'
import numpy as np
import albumentations as A
from albumentations.core.transforms_interface import DualTransform
from albumentations.augmentations.geometric import functional as fgeometric
from packaging import version
ALBU_VERSION = version.parse(A.__version__)
IS_ALBU_NEW_VERSION = ALBU_VERSION >= version.parse('1.4.15')
# Custom resize transformation mimicking Imgaug's behavior with scaling
class ImgaugLikeResize(DualTransform):
def __init__(self, scale_range=(0.5, 3.0), interpolation=1, p=1.0):
super(ImgaugLikeResize, self).__init__(p)
self.scale_range = scale_range
self.interpolation = interpolation
# Resize the image based on a randomly chosen scale within the scale range
def apply(self, img, scale=1.0, **params):
height, width = img.shape[:2]
new_height = int(height * scale)
new_width = int(width * scale)
if IS_ALBU_NEW_VERSION:
return fgeometric.resize(img, (new_height, new_width),
interpolation=self.interpolation)
return fgeometric.resize(img,
new_height,
new_width,
interpolation=self.interpolation)
# Apply the same scaling transformation to keypoints (e.g., polygon points)
def apply_to_keypoints(self, keypoints, scale=1.0, **params):
return np.array([(x * scale, y * scale) + tuple(rest)
for x, y, *rest in keypoints])
# Get random scale parameter within the specified range
def get_params(self):
scale = np.random.uniform(self.scale_range[0], self.scale_range[1])
return {'scale': scale}
# Builder class to translate custom augmenter arguments into Albumentations-compatible format
class AugmenterBuilder(object):
def __init__(self):
# Map common Imgaug transformations to equivalent Albumentations transforms
self.imgaug_to_albu = {
'Fliplr': 'HorizontalFlip',
'Flipud': 'VerticalFlip',
'Affine': 'Affine',
# Additional mappings can be added here if needed
}
# Recursive method to construct augmentation pipeline based on provided arguments
def build(self, args, root=True):
if args is None or len(args) == 0:
return None
elif isinstance(args, list):
# Build the full augmentation sequence if it's a root-level call
if root:
sequence = [self.build(value, root=False) for value in args]
return A.Compose(
sequence,
keypoint_params=A.KeypointParams(format='xy',
remove_invisible=False),
)
else:
# Build individual augmenters for nested arguments
augmenter_type = args[0]
augmenter_args = args[1] if len(args) > 1 else {}
augmenter_args_mapped = self.map_arguments(
augmenter_type, augmenter_args)
augmenter_type_mapped = self.imgaug_to_albu.get(
augmenter_type, augmenter_type)
if augmenter_type_mapped == 'Resize':
return ImgaugLikeResize(**augmenter_args_mapped)
else:
cls = getattr(A, augmenter_type_mapped)
return cls(
**{
k: self.to_tuple_if_list(v)
for k, v in augmenter_args_mapped.items()
})
elif isinstance(args, dict):
# Process individual transformation specified as dictionary
augmenter_type = args['type']
augmenter_args = args.get('args', {})
augmenter_args_mapped = self.map_arguments(augmenter_type,
augmenter_args)
augmenter_type_mapped = self.imgaug_to_albu.get(
augmenter_type, augmenter_type)
if augmenter_type_mapped == 'Resize':
return ImgaugLikeResize(**augmenter_args_mapped)
else:
cls = getattr(A, augmenter_type_mapped)
return cls(
**{
k: self.to_tuple_if_list(v)
for k, v in augmenter_args_mapped.items()
})
else:
raise RuntimeError('Unknown augmenter arg: ' + str(args))
# Map arguments to expected format for each augmenter type
def map_arguments(self, augmenter_type, augmenter_args):
augmenter_args = augmenter_args.copy(
) # Avoid modifying the original arguments
if augmenter_type == 'Resize':
# Ensure size is a valid 2-element list or tuple
size = augmenter_args.get('size')
if size:
if not isinstance(size, (list, tuple)) or len(size) != 2:
raise ValueError(
f"'size' must be a list or tuple of two numbers, but got {size}"
)
min_scale, max_scale = size
return {
'scale_range': (min_scale, max_scale),
'interpolation': 1, # Linear interpolation
'p': 1.0,
}
else:
return {
'scale_range': (1.0, 1.0),
'interpolation': 1,
'p': 1.0
}
elif augmenter_type == 'Affine':
# Map rotation to a tuple and ensure p=1.0 to apply transformation
rotate = augmenter_args.get('rotate', 0)
if isinstance(rotate, list):
rotate = tuple(rotate)
elif isinstance(rotate, (int, float)):
rotate = (float(rotate), float(rotate))
augmenter_args['rotate'] = rotate
augmenter_args['p'] = 1.0
return augmenter_args
else:
# For other augmenters, ensure 'p' probability is specified
p = augmenter_args.get('p', 1.0)
augmenter_args['p'] = p
return augmenter_args
# Convert lists to tuples for Albumentations compatibility
def to_tuple_if_list(self, obj):
if isinstance(obj, list):
return tuple(obj)
return obj
# Wrapper class for image and polygon transformations using Imgaug-style augmentation
class IaaAugment:
def __init__(self, augmenter_args=None, **kwargs):
if augmenter_args is None:
# Default augmenters if none are specified
augmenter_args = [
{
'type': 'Fliplr',
'args': {
'p': 0.5
}
},
{
'type': 'Affine',
'args': {
'rotate': [-10, 10]
}
},
{
'type': 'Resize',
'args': {
'size': [0.5, 3]
}
},
]
self.augmenter = AugmenterBuilder().build(augmenter_args)
# Apply the augmentations to image and polygon data
def __call__(self, data):
image = data['image']
if self.augmenter:
# Flatten polygons to individual keypoints for transformation
keypoints = []
keypoints_lengths = []
for poly in data['polys']:
keypoints.extend([tuple(point) for point in poly])
keypoints_lengths.append(len(poly))
# Apply the augmentation pipeline to image and keypoints
transformed = self.augmenter(image=image, keypoints=keypoints)
data['image'] = transformed['image']
# Extract transformed keypoints and reconstruct polygon structures
transformed_keypoints = transformed['keypoints']
# Reassemble polygons from transformed keypoints
new_polys = []
idx = 0
for length in keypoints_lengths:
new_poly = transformed_keypoints[idx:idx + length]
new_polys.append(np.array([kp[:2] for kp in new_poly]))
idx += length
data['polys'] = np.array(new_polys)
return data
|