Lower-level image manipulation

It is generally far better to operate on whole tensors rather than individual pixels, because reading and writing per-pixel or even per-line data is very slow with native Python. But you can make it work in a pinch. If you need actual performance, take a look at Taichi Lang or Numba.

Source image

Line sweep

import torch
import time
from tinycio import ColorImage
from tinycio.util import apply_hue_oklab

# Each vertical line of the image hue shifted as a tensor.
res = []
im_in = ColorImage.load('../doc/images/examples_ll/horizon.png', 'SRGB')
start = time.time()
im_in = im_in.to_color_space('OKLAB')
_, H, W = im_in.size()
for x in range(W): 
	res.append(apply_hue_oklab(im_in[:, :, x:x+1], x / W * 2. - 1))
im_out = ColorImage(torch.cat(res, dim = 2), 'OKLAB')
im_out = im_out.to_color_space('SRGB').clamp(0., 1.)
end = time.time()
im_out.save('../out/horizon_hue_sweep.png')
print(f'Code execution: {end - start} seconds')

Pixel sweep

import torch
import time
import numpy as np
from tinycio import Color, ColorImage

# Keying into PyTorch tensors for each pixel is prohibitively expensive.
# We can instead hand it over to NumPy - still slow, but relatively tolerable.
im_in = ColorImage.load('../doc/images/examples_ll/horizon.png', color_space='SRGB')
start = time.time()
im_in = im_in.to_color_space('SRGB_LIN').numpy()
C, H, W = im_in.shape
for y in range(H):
    for x in range(W):
        col = Color(im_in[:, y, x])
        col.r *= 1. - (x / W)
        col.g *= 1. - (y / H)
        im_in[:, y, x] = col.rgb
im_out = ColorImage(im_in, 'SRGB_LIN').to_color_space('SRGB').clamp(0., 1.)
end = time.time()
im_out.save('../out/horizon_rg_sweep.png')
print(f'Code execution: {end - start} seconds')

Baseline

import torch
import time
from tinycio import ColorImage

im_in = ColorImage.load('../doc/images/examples_ll/horizon.png', color_space='SRGB')
start = time.time()
im_in = im_in.to_color_space('SRGB_LIN')
C, H, W = im_in.shape
xw = torch.linspace(start=1., end=0., steps=W).unsqueeze(0).repeat(H, 1)
yh = torch.linspace(start=1., end=0., steps=H).unsqueeze(-1).repeat(1, W)
im_in[0,...] *= xw
im_in[1,...] *= yh
im_out = ColorImage(im_in, 'SRGB_LIN').to_color_space('SRGB').clamp(0., 1.)
end = time.time()
im_out.save('../out/horizon_rg_torch.png')
print(f'Code execution: {end - start} seconds')

Same operation, but with PyTorch functions: ~0.04s