diff --git a/Augmentor/Operations.py b/Augmentor/Operations.py index d94623f..d498c8f 100644 --- a/Augmentor/Operations.py +++ b/Augmentor/Operations.py @@ -634,7 +634,15 @@ class RotateStandard(Operation): .. seealso:: For 90 degree rotations, see the :class:`Rotate` class. """ - def __init__(self, probability, max_left_rotation, max_right_rotation, expand=False, fillcolor=None): + def __init__( + self, + probability, + max_left_rotation, + max_right_rotation, + expand=False, + fillcolor=None, + mask_fillcolor=None, + ): """ Documentation to appear. """ @@ -643,6 +651,7 @@ def __init__(self, probability, max_left_rotation, max_right_rotation, expand=Fa self.max_right_rotation = abs(max_right_rotation) # Ensure always positive self.expand = expand self.fillcolor = fillcolor + self.mask_fillcolor = mask_fillcolor def perform_operation(self, images): """ @@ -652,7 +661,6 @@ def perform_operation(self, images): :return: The transformed image(s) as a list of object(s) of type PIL.Image. """ - random_left = random.randint(self.max_left_rotation, 0) random_right = random.randint(0, self.max_right_rotation) @@ -665,13 +673,23 @@ def perform_operation(self, images): elif left_or_right == 1: rotation = random_right - def do(image): - return image.rotate(rotation, expand=self.expand, resample=Image.BICUBIC, fillcolor=self.fillcolor) + def do(image, fillcolor): + return image.rotate( + rotation, + expand=self.expand, + resample=Image.BICUBIC, + fillcolor=fillcolor, + ) augmented_images = [] - for image in images: - augmented_images.append(do(image)) + for i, image in enumerate(images): + # First image is the original - filling it with self.fillcolor + # All of the following images are masks - filling them with `self.mask_fillcolor` + if i == 0: + augmented_images.append(do(image, self.fillcolor)) + else: + augmented_images.append(do(image, self.mask_fillcolor)) return augmented_images diff --git a/Augmentor/Pipeline.py b/Augmentor/Pipeline.py index c7a5aa0..cf8a772 100644 --- a/Augmentor/Pipeline.py +++ b/Augmentor/Pipeline.py @@ -962,7 +962,15 @@ def rotate(self, probability, max_left_rotation, max_right_rotation): self.add_operation(RotateRange(probability=probability, max_left_rotation=ceil(max_left_rotation), max_right_rotation=ceil(max_right_rotation))) - def rotate_without_crop(self, probability, max_left_rotation, max_right_rotation, expand=False, fillcolor=None): + def rotate_without_crop( + self, + probability, + max_left_rotation, + max_right_rotation, + expand=False, + fillcolor=None, + mask_fillcolor=None, + ): """ Rotate an image without automatically cropping. @@ -987,11 +995,15 @@ def rotate_without_crop(self, probability, max_left_rotation, max_right_rotation of the input. Default value is None (points filled with black color). For example, in case of RGB color scheme simply use `(r, g, b)` tuple of int numbers. + :param mask_fillcolor: Specify color to fill points outside the boundaries + of the ground truth. Default value is None (points filled with black color). + For example, in case of RGB color scheme simply use `(r, g, b)` tuple + of int numbers. :return: None """ self.add_operation(RotateStandard(probability=probability, max_left_rotation=ceil(max_left_rotation), max_right_rotation=ceil(max_right_rotation), expand=expand, - fillcolor=fillcolor)) + fillcolor=fillcolor, mask_fillcolor=mask_fillcolor)) def flip_top_bottom(self, probability): """