Add split_and_shuffle augmentation to AugLy (#255)

Summary:
Pull Request resolved: #255

Adding a split and shuffle attack :)

augmentation description: Splits the image into a grid of tiles (determined by n_columns and n_rows) and shuffles the tiles randomly. The resulting image is the concatenation of the shuffled tiles into the same grid format (resulting in an image of the same size)

Reviewed By: vitoralbiero

Differential Revision: D68475586

Author: jbitton

augly/image/__init__.py

@@ -44,6 +44,7 @@
     sharpen,
     shuffle_pixels,
     skew,
+    split_and_shuffle,
     vflip,
 )
 from augly.image.helpers import aug_np_wrapper
@@ -126,6 +127,7 @@
     Sharpen,
     ShufflePixels,
     Skew,
+    SplitAndShuffle,
     VFlip,
 )
 
@@ -174,6 +176,7 @@
     "Sharpen",
     "ShufflePixels",
     "Skew",
+    "SplitAndShuffle",
     "VFlip",
     "apply_lambda",
     "apply_pil_filter",
@@ -211,6 +214,7 @@
     "sharpen",
     "shuffle_pixels",
     "skew",
+    "split_and_shuffle",
     "vflip",
     "apply_lambda_intensity",
     "apply_pil_filter_intensity",

augly/image/functional.py

@@ -13,6 +13,7 @@
 import pickle
 import random
 from copy import deepcopy
+from itertools import product
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -2549,6 +2550,92 @@ def skew(
     return imutils.ret_and_save_image(aug_image, output_path, src_mode)
 
 
+def split_and_shuffle(
+    image: Union[str, Image.Image],
+    output_path: Optional[str] = None,
+    n_columns: int = 3,
+    n_rows: int = 3,
+    seed: int = 10,
+    metadata: Optional[List[Dict[str, Any]]] = None,
+    bboxes: Optional[List[Tuple]] = None,
+    bbox_format: Optional[str] = None,
+) -> Image.Image:
+    """
+    Splits the image into a grid of tiles (determined by n_columns and n_rows) and
+    shuffles the tiles randomly. The resulting image is the concatenation of the
+    shuffled tiles into the same grid format (resulting in an image of the same size)
+
+    @param image: the path to an image or a variable of type PIL.Image.Image
+        to be augmented
+
+    @param output_path: the path in which the resulting image will be stored.
+        If None, the resulting PIL Image will still be returned
+
+    @param n_columns: number of columns to split the image into
+
+    @param n_rows: number of rows to split the image into
+
+    @param seed: seed for numpy random generator to select random order for shuffling
+
+    @param metadata: if set to be a list, metadata about the function execution
+        including its name, the source & dest width, height, etc. will be appended
+        to the inputted list. If set to None, no metadata will be appended or returned
+
+    @param bboxes: a list of bounding boxes can be passed in here if desired. If
+        provided, this list will be modified in place such that each bounding box is
+        transformed according to this function
+
+    @param bbox_format: signifies what bounding box format was used in `bboxes`. Must
+        specify `bbox_format` if `bboxes` is provided. Supported bbox_format values are
+        "pascal_voc", "pascal_voc_norm", "coco", and "yolo"
+
+    @returns: the augmented PIL Image
+    """
+    np.random.seed(seed)
+
+    image = imutils.validate_and_load_image(image)
+
+    assert n_columns > 0, "Expected 'n_columns' to be a positive integer"
+    assert n_rows > 0, "Expected 'n_rows' to be a positive integer"
+
+    func_kwargs = imutils.get_func_kwargs(metadata, locals())
+    src_mode = image.mode
+
+    width, height = image.size
+    width_per_tile = width // n_columns
+    height_per_tile = height // n_rows
+
+    grid = product(
+        range(0, height - height % height_per_tile, height_per_tile),
+        range(0, width - width % width_per_tile, width_per_tile),
+    )
+
+    sub_images = []
+    for y0, x0 in grid:
+        bbox = (x0, y0, x0 + width_per_tile, y0 + height_per_tile)
+        sub_images.append(image.crop(bbox))
+
+    if len(sub_images) == 2:
+        sub_images[0], sub_images[1] = sub_images[1], sub_images[0]
+    else:
+        np.random.shuffle(sub_images)
+
+    aug_image = Image.new("RGB", (width_per_tile * n_columns, height_per_tile * n_rows))
+    for i, sub_image in enumerate(sub_images):
+        x = i % n_columns
+        y = i // n_columns
+        aug_image.paste(sub_image, (x * width_per_tile, y * height_per_tile))
+
+    imutils.get_metadata(
+        metadata=metadata,
+        function_name="split_and_shuffle",
+        aug_image=aug_image,
+        **func_kwargs,
+    )
+
+    return imutils.ret_and_save_image(aug_image, output_path, src_mode)
+
+
 def vflip(
     image: Union[str, Image.Image],
     output_path: Optional[str] = None,

augly/image/intensity.py

@@ -315,6 +315,13 @@ def skew_intensity(skew_factor: float, **kwargs) -> float:
     return min((abs(skew_factor) / max_skew_factor) * 100.0, 100.0)
 
 
+def split_and_shuffle_intensity(n_columns: int, n_rows: int, **kwargs) -> float:
+    assert n_columns > 0, "Expected 'n_columns' to be a positive integer"
+    assert n_rows > 0, "Expected 'n_rows' to be a positive integer"
+
+    return min((1 - (1 / (n_columns * n_rows))) * 100.0, 100.0)
+
+
 def vflip_intensity(**kwargs) -> float:
     return 100.0
 

augly/image/transforms.py

@@ -2207,6 +2207,64 @@ def apply_transform(
         )
 
 
+class SplitAndShuffle(BaseTransform):
+    def __init__(
+        self, n_columns: int = 3, n_rows: int = 3, seed: int = 10, p: float = 1.0
+    ):
+        """
+        @param n_columns: number of columns to split the image into
+
+        @param n_rows: number of rows to split the image into
+
+        @param seed: seed for numpy random generator to select random order
+            for shuffling
+
+        @param p: the probability of the transform being applied; default value is 1.0
+        """
+        super().__init__(p)
+        self.n_columns = n_columns
+        self.n_rows = n_rows
+        self.seed = seed
+
+    def apply_transform(
+        self,
+        image: Image.Image,
+        metadata: Optional[List[Dict[str, Any]]] = None,
+        bboxes: Optional[List[Tuple]] = None,
+        bbox_format: Optional[str] = None,
+    ) -> Image.Image:
+        """
+        Splits the image into a grid of tiles (determined by n_columns and n_rows) and
+        shuffles the tiles randomly. The resulting image is the concatenation of the
+        shuffled tiles into the same grid format (resulting in an image of the same size)
+
+        @param image: PIL Image to be augmented
+
+        @param metadata: if set to be a list, metadata about the function execution
+            including its name, the source & dest width, height, etc. will be appended to
+            the inputted list. If set to None, no metadata will be appended or returned
+
+        @param bboxes: a list of bounding boxes can be passed in here if desired. If
+            provided, this list will be modified in place such that each bounding box is
+            transformed according to this function
+
+        @param bbox_format: signifies what bounding box format was used in `bboxes`. Must
+            specify `bbox_format` if `bboxes` is provided. Supported bbox_format values
+            are "pascal_voc", "pascal_voc_norm", "coco", and "yolo"
+
+        @returns: Augmented PIL Image
+        """
+        return F.split_and_shuffle(
+            image,
+            n_columns=self.n_columns,
+            n_rows=self.n_rows,
+            seed=self.seed,
+            metadata=metadata,
+            bboxes=bboxes,
+            bbox_format=bbox_format,
+        )
+
+
 class VFlip(BaseTransform):
     def apply_transform(
         self,

augly/tests/assets/expected_metadata/image_tests/expected_metadata.json

@@ -741,6 +741,23 @@
             "src_width": 1920
         }
     ],
+    "split_and_shuffle": [
+        {
+            "bbox_format": "yolo",
+            "dst_bboxes": [[0.5, 0.5, 0.25, 0.75]],
+            "dst_height": 1080,
+            "dst_width": 1920,
+            "intensity": 91.66666666666666,
+            "name": "split_and_shuffle",
+            "n_columns": 3,
+            "n_rows": 4,
+            "output_path": null,
+            "seed": 10,
+            "src_bboxes": [[0.5, 0.5, 0.25, 0.75]],
+            "src_height": 1080,
+            "src_width": 1920
+        }
+    ],
     "vflip": [
         {
             "bbox_format": "yolo",

augly/tests/assets/image/dfdc_expected_output/test_split_and_shuffle.png

@@ -138,6 +138,9 @@ def test_shuffle_pixels(self):
     def test_skew(self):
         self.evaluate_function(imaugs.skew)
 
+    def test_split_and_shuffle(self):
+        self.evaluate_function(imaugs.split_and_shuffle, n_columns=3, n_rows=4)
+
     def test_vflip(self):
         self.evaluate_function(imaugs.vflip)
 

augly/tests/image_tests/functional_unit_test.py

@@ -220,6 +220,11 @@ def test_ShufflePixels(self):
     def test_Skew(self):
         self.evaluate_class(imaugs.Skew(), fname="skew")
 
+    def test_SplitAndShuffle(self):
+        self.evaluate_class(
+            imaugs.SplitAndShuffle(n_columns=3, n_rows=4), fname="split_and_shuffle"
+        )
+
     def test_VFlip(self):
         self.evaluate_class(imaugs.VFlip(), fname="vflip")