einops_rearrange_reduce_repeat_image_tensor_manipulation.py

1import numpy as np
2from einops import rearrange, reduce, repeat
3
4# Create a dummy image: batch of 1, 3 channels (RGB), 100x100 pixels
5images = np.random.randn(1, 3, 100, 100)
6
7# 1. Rearrange: Transpose axes (B, C, H, W) -> (B, H, W, C)
8rearranged = rearrange(images, 'b c h w -> b h w c')
9
10# 2. Reduce: Compute the mean across the channel dimension (grayscale)
11grayscale = reduce(images, 'b c h w -> b h w', 'mean')
12
13# 3. Repeat: Stack the grayscale image 3 times to restore channel dimension
14stacked = repeat(grayscale, 'b h w -> b (repeat c) h w', repeat=3)
15
16# 4. Composite: Rearrange and flatten height/width into a single dimension
17flattened = rearrange(images, 'b c h w -> b c (h w)')
18
19print(f"Original shape:  {images.shape}")
20print(f"Rearranged shape: {rearranged.shape}")
21print(f"Grayscale shape:  {grayscale.shape}")
22print(f"Stacked shape:    {stacked.shape}")
23print(f"Flattened shape:  {flattened.shape}")