Skip to content

Commit 567c710

Browse files
Add size limits to RandomScale transform (#462)
1 parent 38db07f commit 567c710

File tree

1 file changed

+18
-3
lines changed

1 file changed

+18
-3
lines changed

mindocr/data/transforms/general_transforms.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,11 +157,19 @@ class RandomScale:
157157
Randomly scales an image and its polygons in a predefined scale range.
158158
Args:
159159
scale_range: (min, max) scale range.
160+
size_limits: (min_side_len, max_side_len) size limits. Default: None.
160161
p: probability of the augmentation being applied to an image.
161162
"""
162163

163-
def __init__(self, scale_range: Union[tuple, list], p: float = 0.5, **kwargs):
164-
self._range = scale_range
164+
def __init__(
165+
self,
166+
scale_range: Union[tuple, list],
167+
size_limits: Union[tuple, list] = None,
168+
p: float = 0.5,
169+
**kwargs,
170+
):
171+
self._range = sorted(scale_range)
172+
self._size_limits = sorted(size_limits) if size_limits else []
165173
self._p = p
166174
assert kwargs.get("is_train", True), ValueError("RandomScale augmentation must be used for training only")
167175

@@ -175,7 +183,14 @@ def __call__(self, data: dict) -> dict:
175183
(polys)
176184
"""
177185
if random.random() < self._p:
178-
scale = np.random.uniform(*self._range)
186+
if self._size_limits:
187+
size = data["image"].shape[:2]
188+
min_scale = max(self._size_limits[0] / size[0], self._size_limits[0] / size[1], self._range[0])
189+
max_scale = min(self._size_limits[1] / size[0], self._size_limits[1] / size[1], self._range[1])
190+
scale = np.random.uniform(min_scale, max_scale)
191+
else:
192+
scale = np.random.uniform(*self._range)
193+
179194
data["image"] = cv2.resize(data["image"], dsize=None, fx=scale, fy=scale)
180195
if "polys" in data:
181196
data["polys"] *= scale

0 commit comments

Comments
 (0)