Skip to content

Instantly share code, notes, and snippets.

@cgranade
Created October 29, 2017 19:51
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save cgranade/4631e1298b02ea5cc90d8f0acfba87f1 to your computer and use it in GitHub Desktop.
Save cgranade/4631e1298b02ea5cc90d8f0acfba87f1 to your computer and use it in GitHub Desktop.
import sys
import numpy as np
from PIL import Image
from PIL.ImageEnhance import Contrast
from PIL.ImageOps import invert
from io import BytesIO
from functools import partial
from toolz import compose
def pil_pipeline_step(fn):
def pipelined_fn(data, *args, **kwargs):
img = Image.open(BytesIO(data))
img = fn(img, *args, **kwargs)
buffer = BytesIO()
img.save(buffer, format='jpeg')
return buffer.getvalue()
return pipelined_fn
@pil_pipeline_step
def flip_image(image):
image.transpose(Image.FLIP_LEFT_RIGHT)
return image
@pil_pipeline_step
def contrast(image, factor):
return Contrast(image).enhance(factor)
invert_colors = pil_pipeline_step(invert)
def randomize_bytes(data, p_error):
data = bytearray(data)
n_errors = np.random.binomial(len(data), p_error)
idxs_errors = np.random.choice(len(data), size=n_errors)
for idx_error in idxs_errors:
data[idx_error] = np.random.randint(0, 256)
return bytes(data)
if __name__ == "__main__":
fname_in = sys.argv[1]
p_error = float(sys.argv[2])
fname_out = sys.argv[3]
with open(fname_in, 'rb') as f:
data = f.read()
pipeline = compose(
flip_image,
invert_colors,
partial(randomize_bytes, p_error=p_error),
invert_colors,
flip_image,
partial(contrast, factor=1.8),
partial(randomize_bytes, p_error=p_error),
)
data = pipeline(data)
with open(fname_out, 'wb') as f:
f.write(data)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment