diff --git a/Augmentor/Operations.py b/Augmentor/Operations.py index 49f9281..bad26a3 100644 --- a/Augmentor/Operations.py +++ b/Augmentor/Operations.py @@ -1154,7 +1154,7 @@ class Shear(Operation): For sample code with image examples see :ref:`shearing`. """ - def __init__(self, probability, max_shear_left, max_shear_right): + def __init__(self, probability, max_shear_left, max_shear_right, directions=['x','y']): """ The shearing is randomised in magnitude, from 0 to the :attr:`max_shear_left` or 0 to :attr:`max_shear_right` where the @@ -1171,6 +1171,7 @@ def __init__(self, probability, max_shear_left, max_shear_right): :type max_shear_right: Integer """ Operation.__init__(self, probability) + self.directions = directions self.max_shear_left = max_shear_left self.max_shear_right = max_shear_right @@ -1227,7 +1228,7 @@ def perform_operation(self, images): # any of the affine transformation matrices, seen here: # https://en.wikipedia.org/wiki/Transformation_matrix#/media/File:2D_affine_transformation_matrix.svg - directions = ["x", "y"] + directions = self.directions # ["x", "y"] direction = random.choice(directions) def do(image): diff --git a/Augmentor/Pipeline.py b/Augmentor/Pipeline.py index c7a5aa0..c8aa35e 100644 --- a/Augmentor/Pipeline.py +++ b/Augmentor/Pipeline.py @@ -1484,7 +1484,7 @@ def skew(self, probability, magnitude=1): skew_type="RANDOM", magnitude=magnitude)) - def shear(self, probability, max_shear_left, max_shear_right): + def shear(self, probability, max_shear_left, max_shear_right, directions=['x','y']): """ Shear the image by a specified number of degrees. @@ -1498,6 +1498,8 @@ def shear(self, probability, max_shear_left, max_shear_right): Cannot be larger than 25 degrees. :param max_shear_right: The max number of degrees to shear to the right. Cannot be larger than 25 degrees. + :param directions: The directions to shear along. + Has to be one of ['x','y'], ['x'], and ['y']. Defaults to ['x','y'] :return: None """ if not 0 < probability <= 1: @@ -1509,7 +1511,8 @@ def shear(self, probability, max_shear_left, max_shear_right): else: self.add_operation(Shear(probability=probability, max_shear_left=max_shear_left, - max_shear_right=max_shear_right)) + max_shear_right=max_shear_right, + directions=directions)) def greyscale(self, probability): """