Support for gravel augmentation (#1365)

* Supports for gravel animations Signed-off-by: Onur Berk Tore <onurberk_t@hotmail.com> * Adds exceptions for tests Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fixes serialization tests Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fixes rest of the tests Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fixes non-rgb warnings. Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fixes the linting errors Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fix more linting Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fix black Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fixes the readme.md Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * Fix the readme Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> * fix: Review updates Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> --------- Signed-off-by: Onur Berk Tore <onurberk_t@hotmail.com> Signed-off-by: delivers-26 <delivers-26@pop-os.localdomain> Co-authored-by: delivers-26 <delivers-26@pop-os.localdomain> Co-authored-by: Mikhail Druzhinin <dipetm@gmail.com>
albumentations-team · Jan 31, 2023 · 7e49d47 · 7e49d47
1 parent 2a1826d
commit 7e49d47
Show file tree

Hide file tree

Showing 6 changed files with 310 additions and 25 deletions.
diff --git a/README.md b/README.md
@@ -153,6 +153,7 @@ Pixel-level transforms will change just an input image and will leave any additi
 - [RandomBrightnessContrast](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomBrightnessContrast)
 - [RandomFog](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomFog)
 - [RandomGamma](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomGamma)
+- [RandomGravel](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomGravel)
 - [RandomRain](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomRain)
 - [RandomShadow](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomShadow)
 - [RandomSnow](https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.RandomSnow)

diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
@@ -25,6 +25,7 @@
     "add_fog",
     "add_rain",
     "add_shadow",
+    "add_gravel",
     "add_snow",
     "add_sun_flare",
     "add_weighted",
@@ -765,6 +766,45 @@ def add_shadow(img, vertices_list):
     return image_rgb
 
 
+@ensure_contiguous
+@preserve_shape
+def add_gravel(img: np.ndarray, gravels: list):
+    """Add gravel to the image.
+
+    From https://github.com/UjjwalSaxena/Automold--Road-Augmentation-Library
+
+    Args:
+        img (numpy.ndarray): image to add gravel to
+        gravels (list): list of gravel parameters. (float, float, float, float):
+            (top-left x, top-left y, bottom-right x, bottom right y)
+
+    Returns:
+        numpy.ndarray:
+    """
+    non_rgb_warning(img)
+    input_dtype = img.dtype
+    needs_float = False
+
+    if input_dtype == np.float32:
+        img = from_float(img, dtype=np.dtype("uint8"))
+        needs_float = True
+    elif input_dtype not in (np.uint8, np.float32):
+        raise ValueError("Unexpected dtype {} for AddGravel augmentation".format(input_dtype))
+
+    image_hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
+
+    for gravel in gravels:
+        y1, y2, x1, x2, sat = gravel
+        image_hls[x1:x2, y1:y2, 1] = sat
+
+    image_rgb = cv2.cvtColor(image_hls, cv2.COLOR_HLS2RGB)
+
+    if needs_float:
+        image_rgb = to_float(image_rgb, max_value=255)
+
+    return image_rgb
+
+
 def invert(img: np.ndarray) -> np.ndarray:
     # Supports all the valid dtypes
     # clips the img to avoid unexpected behaviour.

diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py
@@ -52,6 +52,7 @@
     "FromFloat",
     "RandomBrightnessContrast",
     "RandomSnow",
+    "RandomGravel",
     "RandomRain",
     "RandomFog",
     "RandomSunFlare",
@@ -369,6 +370,117 @@ def get_transform_init_args_names(self):
         return ("snow_point_lower", "snow_point_upper", "brightness_coeff")
 
 
+class RandomGravel(ImageOnlyTransform):
+    """Add gravels.
+
+    From https://github.com/UjjwalSaxena/Automold--Road-Augmentation-Library
+
+    Args:
+        gravel_roi (float, float, float, float): (top-left x, top-left y,
+            bottom-right x, bottom right y). Should be in [0, 1] range
+        number_of_patches (int): no. of gravel patches required
+
+    Targets:
+        image
+
+    Image types:
+        uint8, float32
+    """
+
+    def __init__(
+        self,
+        gravel_roi: tuple = (0.1, 0.4, 0.9, 0.9),
+        number_of_patches: int = 2,
+        always_apply: bool = False,
+        p: float = 0.5,
+    ):
+        super(RandomGravel, self).__init__(always_apply, p)
+
+        (gravel_lower_x, gravel_lower_y, gravel_upper_x, gravel_upper_y) = gravel_roi
+
+        if not 0 <= gravel_lower_x < gravel_upper_x <= 1 or not 0 <= gravel_lower_y < gravel_upper_y <= 1:
+            raise ValueError("Invalid gravel_roi. Got: %s." % gravel_roi)
+        if number_of_patches < 1:
+            raise ValueError("Invalid gravel number_of_patches. Got: %s." % number_of_patches)
+
+        self.gravel_roi = gravel_roi
+        self.number_of_patches = number_of_patches
+
+    def generate_gravel_patch(self, rectangular_roi):
+        x1, y1, x2, y2 = rectangular_roi
+        gravels = []
+        area = abs((x2 - x1) * (y2 - y1))
+        count = area // 10
+        gravels = np.empty([count, 2], dtype=np.int64)
+        gravels[:, 0] = random_utils.randint(x1, x2, count)
+        gravels[:, 1] = random_utils.randint(y1, y2, count)
+        return gravels
+
+    def apply(self, image, gravels_infos=(), **params):
+        return F.add_gravel(image, gravels_infos)
+
+    @property
+    def targets_as_params(self):
+        return ["image"]
+
+    def get_params_dependent_on_targets(self, params):
+        img = params["image"]
+        height, width = img.shape[:2]
+
+        x_min, y_min, x_max, y_max = self.gravel_roi
+        x_min = int(x_min * width)
+        x_max = int(x_max * width)
+        y_min = int(y_min * height)
+        y_max = int(y_max * height)
+
+        max_height = 200
+        max_width = 30
+
+        rectangular_rois = np.zeros([self.number_of_patches, 4], dtype=np.int64)
+        xx1 = random_utils.randint(x_min + 1, x_max, self.number_of_patches)  # xmax
+        xx2 = random_utils.randint(x_min, xx1)  # xmin
+        yy1 = random_utils.randint(y_min + 1, y_max, self.number_of_patches)  # ymax
+        yy2 = random_utils.randint(y_min, yy1)  # ymin
+
+        rectangular_rois[:, 0] = xx2
+        rectangular_rois[:, 1] = yy2
+        rectangular_rois[:, 2] = [min(tup) for tup in zip(xx1, xx2 + max_height)]
+        rectangular_rois[:, 3] = [min(tup) for tup in zip(yy1, yy2 + max_width)]
+
+        minx = []
+        maxx = []
+        miny = []
+        maxy = []
+        val = []
+        for roi in rectangular_rois:
+            gravels = self.generate_gravel_patch(roi)
+            x = gravels[:, 0]
+            y = gravels[:, 1]
+            r = random_utils.randint(1, 4, len(gravels))
+            sat = random_utils.randint(0, 255, len(gravels))
+            miny.append(np.maximum(y - r, 0))
+            maxy.append(np.minimum(y + r, y))
+            minx.append(np.maximum(x - r, 0))
+            maxx.append(np.minimum(x + r, x))
+            val.append(sat)
+
+        return {
+            "gravels_infos": np.stack(
+                [
+                    np.concatenate(miny),
+                    np.concatenate(maxy),
+                    np.concatenate(minx),
+                    np.concatenate(maxx),
+                    np.concatenate(val),
+                ],
+                1,
+            )
+        }
+
+    def get_transform_init_args_names(self):
+        return {"gravel_roi": self.gravel_roi, "number_of_patches": self.number_of_patches}
+
+
 class RandomRain(ImageOnlyTransform):
     """Adds rain effects.
 

diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py
@@ -261,6 +261,7 @@ def test_augmentations_wont_change_float_input(augmentation_cls, params, float_i
             A.PadIfNeeded,
             A.RGBShift,
             A.RandomFog,
+            A.RandomGravel,
             A.RandomRain,
             A.RandomScale,
             A.RandomShadow,
@@ -417,6 +418,7 @@ def test_mask_fill_value(augmentation_cls, params):
             A.RGBShift,
             A.RandomCropNearBBox,
             A.RandomFog,
+            A.RandomGravel,
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
@@ -478,6 +480,7 @@ def test_multichannel_image_augmentations(augmentation_cls, params):
             A.RGBShift,
             A.RandomCropNearBBox,
             A.RandomFog,
+            A.RandomGravel,
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
@@ -532,6 +535,7 @@ def test_float_multichannel_image_augmentations(augmentation_cls, params):
             A.RGBShift,
             A.RandomCropNearBBox,
             A.RandomFog,
+            A.RandomGravel,
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
@@ -588,6 +592,7 @@ def test_multichannel_image_augmentations_diff_channels(augmentation_cls, params
             A.RGBShift,
             A.RandomCropNearBBox,
             A.RandomFog,
+            A.RandomGravel,
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,