A script to randomize the pixels in an image, so long as they're within a mask.

mask_scramble.py

#!/usr/bin/env python
import sys
import argparse

from PIL import Image
import numpy as np


def scramble_ary(img_ar, mask_ar):
    """
    Randomly scamble the pixels in img_ar for which mask_ar is true.

    img_ar should be NxM (probably 3 or 4 -- rgb or rgba)
    mask_ar should be Nx1 (boolean)

    Returns an array shaped like img_ar.

    """
    idxs_orig = np.arange(img_ar.shape[0])
    idxs_new = np.copy(idxs_orig)

    shuffle_orig = np.where(mask_ar)[0]
    shuffle_new = np.copy(shuffle_orig)
    np.random.shuffle(shuffle_new)

    idxs_new[shuffle_new] = idxs_orig[shuffle_orig]
    img_out = np.copy(img_ar)
    img_out[idxs_new] = img_ar[idxs_orig]
    return img_out


def make_flat_mask(mask_file, expected_shape):
    """
    Make a 1-dimensional boolean mask -- True wherever the image is non-black.

    Raise a RuntimeError if the mask's shape is not expected_shape.
    """
    # Open the imagem convert it to greyscale, and make an array.
    mask_ar = np.array(Image.open(mask_file).convert("L"))
    if not np.array_equal(mask_ar.shape[0:2], expected_shape[0:2]):
        raise RuntimeError("mask must be the same shape as in_file")
    mask_ar = mask_ar.reshape(expected_shape[0]*expected_shape[1])
    return (mask_ar > 0)


def scramble_image(in_file, mask_file=None):
    img_ar = np.array(Image.open(in_file))
    num_pix = img_ar.shape[0]*img_ar.shape[1]
    img_flat = img_ar.reshape(num_pix, -1)
    mask = np.ones(num_pix, dtype=bool)
    if mask_file is not None:
        mask = make_flat_mask(mask_file, img_ar.shape)
    scrambled = scramble_ary(img_flat, mask)
    return Image.fromarray(scrambled.reshape(img_ar.shape))


def main(in_file, out_file, mask=None):
    out_img = scramble_image(in_file, mask)
    out_img.save(out_file)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Randomly shuffle the pixels in an image.")
    parser.add_argument(
        '--mask', nargs='?', default=None,
        help="a mask, should be nonzero for pixels you want scrambled")
    parser.add_argument('in_file')
    parser.add_argument('out_file')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    sys.exit(main(args.in_file, args.out_file, args.mask))