upgrade to numpy 2.0 and remove imgaug

Author: GreatV

ppocr/data/imaug/iaa_augment.py

@@ -19,37 +19,88 @@
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
+import os
 
 import numpy as np
-import imgaug
-import imgaug.augmenters as iaa
 
+os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
+import albumentations as A
 
-class AugmenterBuilder(object):
+
+class AugmenterBuilder:
     def __init__(self):
         pass
 
-    def build(self, args, root=True):
-        if args is None or len(args) == 0:
+    def build(self, args):
+        if not args:
             return None
         elif isinstance(args, list):
-            if root:
-                sequence = [self.build(value, root=False) for value in args]
-                return iaa.Sequential(sequence)
+            # Recursively build transforms from the list
+            transforms = [self.build(value) for value in args if self.build(value)]
+            return A.Compose(transforms, keypoint_params=A.KeypointParams(format="xy"))
+        elif isinstance(args, dict):
+            # Get the transform type and its arguments
+            transform_type = args.get("type")
+            transform_args = args.get("args", {})
+            # Map the transform type to the corresponding function
+            transform_func = self._get_transform_function(
+                transform_type, transform_args
+            )
+            if transform_func:
+                return transform_func
             else:
-                return getattr(iaa, args[0])(
-                    *[self.to_tuple_if_list(a) for a in args[1:]]
+                raise NotImplementedError(
+                    f"Transform {transform_type} not implemented."
                 )
-        elif isinstance(args, dict):
-            cls = getattr(iaa, args["type"])
-            return cls(**{k: self.to_tuple_if_list(v) for k, v in args["args"].items()})
         else:
-            raise RuntimeError("unknown augmenter arg: " + str(args))
+            raise RuntimeError(f"Unknown augmenter arg: {args}")
+
+    def _get_transform_function(self, transform_type, transform_args):
+        # Define mapping from transform types to functions
+        transform_mapping = {
+            "Fliplr": self._build_horizontal_flip,
+            "Affine": self._build_affine,
+            "Resize": self._build_resize,
+        }
+        func = transform_mapping.get(transform_type)
+        if func:
+            return func(transform_args)
+        else:
+            return None
 
-    def to_tuple_if_list(self, obj):
-        if isinstance(obj, list):
-            return tuple(obj)
-        return obj
+    def _build_horizontal_flip(self, transform_args):
+        p = transform_args.get("p", 0.5)
+        return A.HorizontalFlip(p=p)
+
+    def _build_affine(self, transform_args):
+        rotate = transform_args.get("rotate")
+        shear = transform_args.get("shear")
+        translate_percent = transform_args.get("translate_percent")
+        affine_args = {"fit_output": True}
+        if rotate is not None:
+            affine_args["rotate"] = (
+                tuple(rotate) if isinstance(rotate, list) else rotate
+            )
+        if shear is not None:
+            affine_args["shear"] = shear
+        if translate_percent is not None:
+            affine_args["translate_percent"] = translate_percent
+        return A.Affine(**affine_args)
+
+    def _build_resize(self, transform_args):
+        size = transform_args.get("size", [1.0, 1.0])
+        if isinstance(size, list) and len(size) == 2:
+            scale_factor = size[0]
+            height = int(scale_factor * 100)
+            width = int(scale_factor * 100)
+            return A.Resize(height=height, width=width)
+        elif isinstance(size, (int, float)):
+            scale_factor = float(size)
+            height = int(scale_factor * 100)
+            width = int(scale_factor * 100)
+            return A.Resize(height=height, width=width)
+        else:
+            raise ValueError("Invalid size parameter for Resize")
 
 
 class IaaAugment:
@@ -58,35 +109,28 @@ def __init__(self, augmenter_args=None, **kwargs):
             augmenter_args = [
                 {"type": "Fliplr", "args": {"p": 0.5}},
                 {"type": "Affine", "args": {"rotate": [-10, 10]}},
-                {"type": "Resize", "args": {"size": [0.5, 3]}},
+                {"type": "Resize", "args": {"size": [0.5, 3.0]}},
             ]
         self.augmenter = AugmenterBuilder().build(augmenter_args)
 
     def __call__(self, data):
         image = data["image"]
-        shape = image.shape
-
+        polys = data["polys"]
+        # Flatten polys to a list of keypoints
+        keypoints = [tuple(point) for poly in polys for point in poly]
         if self.augmenter:
-            aug = self.augmenter.to_deterministic()
-            data["image"] = aug.augment_image(image)
-            data = self.may_augment_annotation(aug, data, shape)
-        return data
-
-    def may_augment_annotation(self, aug, data, shape):
-        if aug is None:
-            return data
-
-        line_polys = []
-        for poly in data["polys"]:
-            new_poly = self.may_augment_poly(aug, shape, poly)
-            line_polys.append(new_poly)
-        data["polys"] = np.array(line_polys)
+            transformed = self.augmenter(image=image, keypoints=keypoints)
+            data["image"] = transformed["image"]
+            # Reconstruct polys from transformed keypoints
+            transformed_keypoints = transformed["keypoints"]
+            new_polys = []
+            idx = 0
+            for poly in polys:
+                num_points = len(poly)
+                new_poly = np.array(
+                    transformed_keypoints[idx : idx + num_points], dtype=np.float32
+                )
+                new_polys.append(new_poly)
+                idx += num_points
+            data["polys"] = new_polys
         return data
-
-    def may_augment_poly(self, aug, img_shape, poly):
-        keypoints = [imgaug.Keypoint(p[0], p[1]) for p in poly]
-        keypoints = aug.augment_keypoints(
-            [imgaug.KeypointsOnImage(keypoints, shape=img_shape)]
-        )[0].keypoints
-        poly = [(p.x, p.y) for p in keypoints]
-        return poly

pyproject.toml

@@ -41,11 +41,10 @@ classifiers = [
 dependencies = [
     "shapely",
     "scikit-image",
-    "imgaug",
     "pyclipper",
     "lmdb",
     "tqdm",
-    "numpy<2.0",
+    "numpy",
     "rapidfuzz",
     "opencv-python",
     "opencv-contrib-python",
@@ -56,7 +55,10 @@ dependencies = [
     "beautifulsoup4",
     "fonttools>=4.24.0",
     "fire>=0.3.0",
-    "requests"
+    "requests",
+    "albumentations==1.4.10",
+    # to be compatible with albumentations
+    "albucore==0.0.13"
 ]
 
 [project.urls]

requirements.txt

@@ -1,10 +1,9 @@
 shapely
 scikit-image
-imgaug
 pyclipper
 lmdb
 tqdm
-numpy<2.0
+numpy
 rapidfuzz
 opencv-python
 opencv-contrib-python

tests/test_iaa_augment.py

@@ -0,0 +1,174 @@
+import os
+import sys
+import pytest
+import numpy as np
+import random
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(current_dir, "..")))
+
+from ppocr.data.imaug.iaa_augment import IaaAugment
+
+# Set a fixed random seed for reproducibility
+np.random.seed(42)
+random.seed(42)
+
+
+# Fixtures for common test inputs
+@pytest.fixture
+def sample_image():
+    # Create a dummy image of size 100x100 with 3 channels
+    return np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
+
+
+@pytest.fixture
+def sample_polys():
+    # Create some dummy polygons
+    polys = [
+        np.array([[10, 10], [20, 10], [20, 20], [10, 20]], dtype=np.float32),
+        np.array([[30, 30], [40, 30], [40, 40], [30, 40]], dtype=np.float32),
+    ]
+    return polys
+
+
+# Helper function to create data dictionary
+def create_data(sample_image, sample_polys):
+    return {
+        "image": sample_image.copy(),
+        "polys": [poly.copy() for poly in sample_polys],
+    }
+
+
+# Test default augmenter (with default augmenter_args)
+def test_iaa_augment_default(sample_image, sample_polys):
+    data = create_data(sample_image, sample_polys)
+    augmenter = IaaAugment()
+    transformed_data = augmenter(data)
+
+    assert isinstance(
+        transformed_data["image"], np.ndarray
+    ), "Image should be a numpy array"
+    assert isinstance(transformed_data["polys"], list), "Polys should be a list"
+    assert transformed_data["image"].ndim == 3, "Image should be 3-dimensional"
+
+    # Check that the polys have been transformed
+    polys_changed = any(
+        not np.allclose(orig_poly, trans_poly)
+        for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
+    )
+    assert polys_changed, "Polygons should have been transformed"
+
+
+# Test augmenter with empty augmenter_args (no augmentation)
+def test_iaa_augment_none(sample_image, sample_polys):
+    data = create_data(sample_image, sample_polys)
+    augmenter = IaaAugment(augmenter_args=[])
+    transformed_data = augmenter(data)
+
+    assert np.array_equal(
+        data["image"], transformed_data["image"]
+    ), "Image should be unchanged"
+    for orig_poly, transformed_poly in zip(data["polys"], transformed_data["polys"]):
+        assert np.array_equal(
+            orig_poly, transformed_poly
+        ), "Polygons should be unchanged"
+
+
+# Parameterize tests to cover multiple augmenter_args scenarios
+@pytest.mark.parametrize(
+    "augmenter_args, expected_shape",
+    [
+        ([], (100, 100, 3)),
+        ([{"type": "Resize", "args": {"size": [0.5, 0.5]}}], (50, 50, 3)),
+        ([{"type": "Resize", "args": {"size": [2.0, 2.0]}}], (200, 200, 3)),
+    ],
+)
+def test_iaa_augment_resize(sample_image, sample_polys, augmenter_args, expected_shape):
+    data = create_data(sample_image, sample_polys)
+    augmenter = IaaAugment(augmenter_args=augmenter_args)
+    transformed_data = augmenter(data)
+
+    assert (
+        transformed_data["image"].shape == expected_shape
+    ), f"Expected image shape {expected_shape}, got {transformed_data['image'].shape}"
+
+
+# Test with custom augmenter_args
+def test_iaa_augment_custom(sample_image, sample_polys):
+    data = create_data(sample_image, sample_polys)
+    augmenter_args = [
+        {"type": "Affine", "args": {"rotate": [45, 45]}},  # Fixed rotation angle
+        {"type": "Resize", "args": {"size": [0.5, 0.5]}},
+    ]
+    augmenter = IaaAugment(augmenter_args=augmenter_args)
+    transformed_data = augmenter(data)
+
+    expected_height = int(sample_image.shape[0] * 0.5)
+    expected_width = int(sample_image.shape[1] * 0.5)
+
+    assert (
+        transformed_data["image"].shape[0] == expected_height
+    ), "Image height should be scaled by 0.5"
+    assert (
+        transformed_data["image"].shape[1] == expected_width
+    ), "Image width should be scaled by 0.5"
+
+    # Check that the polys have been transformed
+    polys_changed = any(
+        not np.allclose(orig_poly, trans_poly)
+        for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
+    )
+    assert polys_changed, "Polygons should have been transformed"
+
+
+# Test unknown transform type raises NotImplementedError
+def test_iaa_augment_unknown_transform():
+    augmenter_args = [{"type": "UnknownTransform", "args": {}}]
+    with pytest.raises(NotImplementedError):
+        IaaAugment(augmenter_args=augmenter_args)
+
+
+# Test invalid size parameter raises ValueError
+def test_iaa_augment_invalid_resize_size():
+    augmenter_args = [{"type": "Resize", "args": {"size": "invalid_size"}}]
+    with pytest.raises(ValueError):
+        IaaAugment(augmenter_args=augmenter_args)
+
+
+# Test that polys are transformed appropriately
+def test_iaa_augment_polys_transformation(sample_image, sample_polys):
+    data = create_data(sample_image, sample_polys)
+    augmenter_args = [
+        {"type": "Affine", "args": {"rotate": [90, 90]}},  # Fixed rotation angle
+    ]
+    augmenter = IaaAugment(augmenter_args=augmenter_args)
+    transformed_data = augmenter(data)
+
+    # Check that the polygons have changed
+    polys_changed = any(
+        not np.allclose(orig_poly, trans_poly)
+        for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
+    )
+    assert polys_changed, "Polygons should have been transformed"
+
+
+# Test with multiple transforms in augmenter_args
+def test_iaa_augment_multiple_transforms(sample_image, sample_polys):
+    augmenter_args = [
+        {"type": "Fliplr", "args": {"p": 1.0}},  # Always flip
+        {"type": "Affine", "args": {"shear": 10}},
+    ]
+    data = create_data(sample_image, sample_polys)
+    augmenter = IaaAugment(augmenter_args=augmenter_args)
+    transformed_data = augmenter(data)
+
+    # Check that the image has been transformed
+    images_different = not np.array_equal(transformed_data["image"], sample_image)
+    assert images_different, "Image should be transformed"
+
+    # Check that the polys have been transformed
+    polys_changed = any(
+        not np.allclose(orig_poly, trans_poly)
+        for orig_poly, trans_poly in zip(sample_polys, transformed_data["polys"])
+    )
+    assert polys_changed, "Polygons should have been transformed"