Ehtesham123/OCR_AEB_Serial_Number

1# Scene Text Recognition Model Hub
2# Copyright 2022 Darwin Bautista
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     https://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16from functools import partial
17
18import imgaug.augmenters as iaa
19import numpy as np
20from PIL import Image, ImageFilter
21
22from timm.data import auto_augment
23
24from strhub.data import aa_overrides
25
26aa_overrides.apply()
27
28_OP_CACHE = {}
29
30
31def _get_op(key, factory):
32    try:
33        op = _OP_CACHE[key]
34    except KeyError:
35        op = factory()
36        _OP_CACHE[key] = op
37    return op
38
39
40def _get_param(level, img, max_dim_factor, min_level=1):
41    max_level = max(min_level, max_dim_factor * max(img.size))
42    return round(min(level, max_level))
43
44
45def gaussian_blur(img, radius, **__):
46    radius = _get_param(radius, img, 0.02)
47    key = 'gaussian_blur_' + str(radius)
48    op = _get_op(key, lambda: ImageFilter.GaussianBlur(radius))
49    return img.filter(op)
50
51
52def motion_blur(img, k, **__):
53    k = _get_param(k, img, 0.08, 3) | 1  # bin to odd values
54    key = 'motion_blur_' + str(k)
55    op = _get_op(key, lambda: iaa.MotionBlur(k))
56    return Image.fromarray(op(image=np.asarray(img)))
57
58
59def gaussian_noise(img, scale, **_):
60    scale = _get_param(scale, img, 0.25) | 1  # bin to odd values
61    key = 'gaussian_noise_' + str(scale)
62    op = _get_op(key, lambda: iaa.AdditiveGaussianNoise(scale=scale))
63    return Image.fromarray(op(image=np.asarray(img)))
64
65
66def poisson_noise(img, lam, **_):
67    lam = _get_param(lam, img, 0.2) | 1  # bin to odd values
68    key = 'poisson_noise_' + str(lam)
69    op = _get_op(key, lambda: iaa.AdditivePoissonNoise(lam))
70    return Image.fromarray(op(image=np.asarray(img)))
71
72
73def _level_to_arg(level, _hparams, max):
74    level = max * level / auto_augment._LEVEL_DENOM
75    return (level,)
76
77
78_RAND_TRANSFORMS = auto_augment._RAND_INCREASING_TRANSFORMS.copy()
79_RAND_TRANSFORMS.remove('SharpnessIncreasing')  # remove, interferes with *blur ops
80_RAND_TRANSFORMS.extend([
81    'GaussianBlur',
82    # 'MotionBlur',
83    # 'GaussianNoise',
84    'PoissonNoise',
85])
86auto_augment.LEVEL_TO_ARG.update({
87    'GaussianBlur': partial(_level_to_arg, max=4),
88    'MotionBlur': partial(_level_to_arg, max=20),
89    'GaussianNoise': partial(_level_to_arg, max=0.1 * 255),
90    'PoissonNoise': partial(_level_to_arg, max=40),
91})
92auto_augment.NAME_TO_OP.update({
93    'GaussianBlur': gaussian_blur,
94    'MotionBlur': motion_blur,
95    'GaussianNoise': gaussian_noise,
96    'PoissonNoise': poisson_noise,
97})
98
99
100def rand_augment_transform(magnitude=5, num_layers=3):
101    # These are tuned for magnitude=5, which means that effective magnitudes are half of these values.
102    hparams = {
103        'rotate_deg': 30,
104        'shear_x_pct': 0.9,
105        'shear_y_pct': 0.2,
106        'translate_x_pct': 0.10,
107        'translate_y_pct': 0.30,
108    }
109    ra_ops = auto_augment.rand_augment_ops(magnitude, hparams=hparams, transforms=_RAND_TRANSFORMS)
110    # Supply weights to disable replacement in random selection (i.e. avoid applying the same op twice)
111    choice_weights = [1.0 / len(ra_ops) for _ in range(len(ra_ops))]
112    return auto_augment.RandAugment(ra_ops, num_layers, choice_weights)
113