import tensorflow as tf
import numpy as np


def image_to_patches(image, patch_size, stride):
    target_width = ((image.shape[1] - patch_size) // stride + 1) * stride + patch_size + 1
    target_height = ((image.shape[0] - patch_size) // stride + 1) * stride + patch_size + 1

    image = np.pad(
        image,
        ((0, target_height - image.shape[0]), (0, target_width - image.shape[1]), (0, 0)),
        mode='constant'
    )
    print(f'Last pixels of the bottom row of the image:\n {image[-1, -5:]}')
    # here, the last row of `image` is all zeros
    batched_image = np.expand_dims(image, 0)
    patches = tf.image.extract_patches(
        images=batched_image,
        sizes=[1, patch_size, patch_size, 1],
        strides=[1, stride, stride, 1],
        rates=[1, 1, 1, 1],
        padding='VALID',
        name=None
    )
    patches = np.array(patches)
    patches = np.resize(patches, (patches.shape[0] * patches.shape[1] * patches.shape[2], -1))
    patches = np.resize(patches, (patches.shape[0], patch_size, patch_size, image.shape[-1]))
    print(f'Last pixels of the bottom row of the last patch:\n {patches[-1, -1, -5:]}')
    return patches


def main():
    image = 255 * np.ones((200, 100, 3))
    _ = image_to_patches(image, patch_size=60, stride=50)


if __name__ == '__main__':
    main()