Add size limits to RandomScale transform

mindocr/data/transforms/general_transforms.py

@@ -157,11 +157,19 @@ class RandomScale:
     Randomly scales an image and its polygons in a predefined scale range.
     Args:
         scale_range: (min, max) scale range.
+        size_limits: (min_side_len, max_side_len) size limits. Default: None.
         p: probability of the augmentation being applied to an image.
     """
 
-    def __init__(self, scale_range: Union[tuple, list], p: float = 0.5, **kwargs):
-        self._range = scale_range
+    def __init__(
+        self,
+        scale_range: Union[tuple, list],
+        size_limits: Union[tuple, list] = None,
+        p: float = 0.5,
+        **kwargs,
+    ):
+        self._range = sorted(scale_range)
+        self._size_limits = sorted(size_limits) if size_limits else []
         self._p = p
         assert kwargs.get("is_train", True), ValueError("RandomScale augmentation must be used for training only")
 
@@ -175,7 +183,14 @@ def __call__(self, data: dict) -> dict:
             (polys)
         """
         if random.random() < self._p:
-            scale = np.random.uniform(*self._range)
+            if self._size_limits:
+                size = data["image"].shape[:2]
+                min_scale = max(self._size_limits[0] / size[0], self._size_limits[0] / size[1], self._range[0])
+                max_scale = min(self._size_limits[1] / size[0], self._size_limits[1] / size[1], self._range[1])
+                scale = np.random.uniform(min_scale, max_scale)
+            else:
+                scale = np.random.uniform(*self._range)
+
             data["image"] = cv2.resize(data["image"], dsize=None, fx=scale, fy=scale)
             if "polys" in data:
                 data["polys"] *= scale