ViT Patchify

Problem

Split an image into non-overlapping patch x patch patches for a Vision Transformer.

Iterate patches in row-major order and return each patch as a flat list of patch * patch pixel values (row-major inside the patch). You may assume both image dimensions are divisible by patch.

Hint

Walk patch starts with range(0, h, patch) and range(0, w, patch).

Test cases

case 1

in: [[[1,2],[3,4]],2]

expected: [[1,2,3,4]]

case 2

in: [[[0,1,2,3],[4,5,6,7],[8,9,10,11],[12,13,14,15]],2]

expected: [[0,1,4,5],[2,3,6,7],[8,9,12,13],[10,11,14,15]]

case 3

in: [[[1,2,3],[4,5,6],[7,8,9]],3]

expected: [[1,2,3,4,5,6,7,8,9]]

case 4

in: [[[1,2],[3,4]],1]

expected: [[1],[2],[3],[4]]

Starter codePython

def patchify(image, patch):
    # Your code here
    pass

4 test cases · run your code against them in the editor