diff --git a/augly/image/__init__.py b/augly/image/__init__.py index e33232c3..4c9ab836 100644 --- a/augly/image/__init__.py +++ b/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", diff --git a/augly/image/functional.py b/augly/image/functional.py index fbea3c01..fbdee6f0 100644 --- a/augly/image/functional.py +++ b/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, diff --git a/augly/image/intensity.py b/augly/image/intensity.py index 56004f74..7ab15b73 100644 --- a/augly/image/intensity.py +++ b/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 diff --git a/augly/image/transforms.py b/augly/image/transforms.py index 623e5f0c..eb39daff 100644 --- a/augly/image/transforms.py +++ b/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, diff --git a/augly/tests/assets/expected_metadata/image_tests/expected_metadata.json b/augly/tests/assets/expected_metadata/image_tests/expected_metadata.json index 28b0afcc..03847a9a 100644 --- a/augly/tests/assets/expected_metadata/image_tests/expected_metadata.json +++ b/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", diff --git a/augly/tests/assets/image/dfdc_expected_output/test_split_and_shuffle.png b/augly/tests/assets/image/dfdc_expected_output/test_split_and_shuffle.png new file mode 100644 index 00000000..f4ade685 Binary files /dev/null and b/augly/tests/assets/image/dfdc_expected_output/test_split_and_shuffle.png differ diff --git a/augly/tests/image_tests/functional_unit_test.py b/augly/tests/image_tests/functional_unit_test.py index 762d62fd..4d346eb9 100644 --- a/augly/tests/image_tests/functional_unit_test.py +++ b/augly/tests/image_tests/functional_unit_test.py @@ -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) diff --git a/augly/tests/image_tests/transforms_unit_test.py b/augly/tests/image_tests/transforms_unit_test.py index 869ed835..df32e211 100644 --- a/augly/tests/image_tests/transforms_unit_test.py +++ b/augly/tests/image_tests/transforms_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")