Easy & Intuitive Image Chopping Into 16x16 Chunks!

[akmayer]

This package has quickly become my favorite for working with tensors. One thing Vision Transformers do is break images into a grid of 16 by 16 chunks and einx makes this an incredibly intuitive and easy one liner in my opinion :). Just wanted to show my appreciation for the great work in this package:

import einx
sampleFrame = frameDataset[18000]
plt.imshow(einx.rearrange("c h w -> h w c", sampleFrame))

chunkedFrame = einx.rearrange("c (h h16) (w w16) -> h w h16 w16 c", sampleFrame, h16=16, w16=16)
heightStrip = chunkedFrame[:, 20]
plt.imshow(einx.rearrange("h h16 w16 c -> (h h16) w16 c", heightStrip))