Split and shuffle

zigsaw.py

def zigsaw(image, num_cuts, random_state=None, **kwargs):
    """Split and shuffle
    Args:
        image: Input image, currently, only RGB, uint8 images are supported.
        num_cuts: number of cuts per side. Image is split `into num_cuts * num_cuts` pieces
        random_state:
        **kwargs:
    Returns:
        Image with reshuffled patches.
    """
    if random_state is None:
        random_state = np.random.RandomState(42)

    height, width, *rest = image.shape

    assert height % num_cuts == 0
    assert width % num_cuts == 0

    block_height = height // num_cuts
    block_width = width // num_cuts
    
    splitted_h = np.swapaxes(image.reshape(num_cuts, block_height, width, *rest), 1, 2)
    splitted = np.swapaxes(splitted_h.reshape(num_cuts * num_cuts, block_width, block_height, *rest), 2, 1)
    
    random_state.shuffle(splitted)

    image_out = np.hstack(np.hstack(splitted.reshape(num_cuts, num_cuts, block_height, block_width, *rest)))

    return image_out