Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unhardcode resample method for Distort operations #226

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Unhardcode resample method for Distort operations #226

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions Augmentor/Operations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@
# except ImportError:
# from io import StringIO

PIL_filters = {
'NEAREST' : Image.NEAREST,
'BILINEAR': Image.BILINEAR,
'BICUBIC' : Image.BICUBIC
}

class Operation(object):
"""
Expand Down Expand Up @@ -1356,7 +1361,7 @@ class Distort(Operation):
"""
This class performs randomised, elastic distortions on images.
"""
def __init__(self, probability, grid_width, grid_height, magnitude):
def __init__(self, probability, grid_width, grid_height, magnitude, resample="BICUBIC"):
"""
As well as the probability, the granularity of the distortions
produced by this class can be controlled using the width and
Expand Down Expand Up @@ -1385,6 +1390,7 @@ def __init__(self, probability, grid_width, grid_height, magnitude):
self.magnitude = abs(magnitude)
# TODO: Implement non-random magnitude.
self.randomise_magnitude = True
self.resample_method = PIL_filters[resample]

def perform_operation(self, images):
"""
Expand Down Expand Up @@ -1486,7 +1492,7 @@ def perform_operation(self, images):

def do(image):

return image.transform(image.size, Image.MESH, generated_mesh, resample=Image.BICUBIC)
return image.transform(image.size, Image.MESH, generated_mesh, resample=self.resample_method)

augmented_images = []

Expand All @@ -1500,7 +1506,7 @@ class GaussianDistortion(Operation):
"""
This class performs randomised, elastic gaussian distortions on images.
"""
def __init__(self, probability, grid_width, grid_height, magnitude, corner, method, mex, mey, sdx, sdy):
def __init__(self, probability, grid_width, grid_height, magnitude, corner, method, mex, mey, sdx, sdy, resample='BICUBIC'):
"""
As well as the probability, the granularity of the distortions
produced by this class can be controlled using the width and
Expand Down Expand Up @@ -1560,6 +1566,7 @@ def __init__(self, probability, grid_width, grid_height, magnitude, corner, meth
self.mey = mey
self.sdx = sdx
self.sdy = sdy
self.resample_method = PIL_filters[resample]

def perform_operation(self, images):
"""
Expand Down Expand Up @@ -1686,7 +1693,7 @@ def do(image):
for i in range(len(dimensions)):
generated_mesh.append([dimensions[i], polygons[i]])

return image.transform(image.size, Image.MESH, generated_mesh, resample=Image.BICUBIC)
return image.transform(image.size, Image.MESH, generated_mesh, resample=self.resample_method)

augmented_images = []

Expand Down