PyTorch Image Transforms by Hand

An image transform is a change made to an image tensor, like a flip, a crop, or a shift. Slicing and flip calls do these jobs directly on the tensor.

A flip on a tiny grid

Flipping dimension 1 mirrors each row. The single bright pixel moves from left to right.

import torch

grid = torch.tensor([[1.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0]])
print(torch.flip(grid, dims=[1]))

A crop is a slice

Keeping rows 1 to 2 and columns 1 to 2 cuts the centre out of a 4 by 4 grid.

import torch

grid = torch.arange(16, dtype=torch.float32).reshape(4, 4)
crop = grid[1:3, 1:3]
print(crop)

A shift with roll

Roll moves every pixel by the shift amount. Pixels pushed off one edge return on the other edge.

import torch

grid = torch.tensor([[1.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0]])
print(torch.roll(grid, shifts=1, dims=1))

Example

import torch

# One square image, built the same way as the shapes dataset.
image = torch.zeros(16, 16)
image[3:11, 2:10] = 1.0


def show(name, img):
    print(name)
    for row in img:
        print("".join("#" if v.item() > 0.5 else "." for v in row))


show("Original:", image)
show("Flip left-right:", torch.flip(image, dims=[1]))
show("Flip up-down:", torch.flip(image, dims=[0]))
show("Shift by roll:", torch.roll(image, shifts=(3, 4), dims=(0, 1)))
show("Centre crop:", image[4:12, 4:12])

Key points

Transform: a change to an image, like flipping, cropping, or shifting it
Use torch.flip to mirror an image along a chosen dimension
Use slicing to crop, by keeping only a range of rows and columns
Use torch.roll to shift an image, with pixels wrapping around the edge
Augmentation: making changed copies of training images so the model sees more variety
Each transform is a plain tensor operation, and libraries wrap these same operations

In this example, one square image starts with its bright block left of centre. The first grid prints the original. The flip on dimension 1 mirrors the block to the right side. The flip on dimension 0 moves the block from the top rows to the lower rows. The roll shifts the block down and right, and no pixels are lost. The centre crop keeps rows and columns 4 to 11, so the grid shrinks to 8 by 8. Four transforms come from three plain operations. The operations are flip, roll, and slicing.

Note

torch.flip returns a new tensor and leaves the original image unchanged

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

# One square image, built the same way as the shapes dataset.
image = torch.zeros(16, 16)
image[3:11, 2:10] = 1.0


def show(name, img):
    print(name)
    for row in img:
        print("".join("#" if v.item() > 0.5 else "." for v in row))


show("Original:", image)
show("Flip left-right:", torch.flip(image, dims=[1]))
show("Flip up-down:", torch.flip(image, dims=[0]))
show("Shift by roll:", torch.roll(image, shifts=(3, 4), dims=(0, 1)))
show("Centre crop:", image[4:12, 4:12])

Exercise

Which operation crops an image tensor?

Slicing a range of rows and columns torch.flip on dimension 0 torch.roll with a shift of 2 Multiplying every pixel by 0