example-0a551db2f9d6.py

 1import numpy as np
 2from numpy.typing import NDArray
 3
 4
 5def patchify(image: NDArray[np.float64], patch_size: int) -> NDArray[np.float64]:
 6    h, w, c = image.shape
 7    p = patch_size
 8    if p <= 0 or h % p or w % p:
 9        raise ValueError("Patch size must divide height and width")
10    grid = image.reshape(h // p, p, w // p, p, c)
11    patches = grid.transpose(0, 2, 1, 3, 4)
12    return patches.reshape((h // p) * (w // p), p * p * c)
13
14
15image = np.arange(224 * 224 * 3, dtype=np.float64).reshape(224, 224, 3)
16patches = patchify(image, patch_size=16)
17
18assert patches.shape == (196, 768)
19np.testing.assert_array_equal(
20    patches[0], image[:16, :16, :].reshape(-1)
21)
22np.testing.assert_array_equal(
23    patches[1], image[:16, 16:32, :].reshape(-1)
24)
25print(patches.shape)  # (196, 768)