Quick pixel manipulation with Pillow and/or NumPy

Question

I'm trying to improve the speed of my image manipulation as it's been too slow for actual use.

What I need to do is apply a complex transformation on the colour of every pixel on an image. The manipulation is basically apply a vector transform like T(r, g, b, a) => (r * x, g * x, b * y, a) or in layman's terms, it's a multiplication of Red and Green values by a constant, a different multiplication for Blue and keep Alpha. But I also need to manipulate it differently if the RGB colour falls under some specific colours, in those cases they must follow a dictionary/transformation table where RGB => newRGB again keeping alpha.

The algorithm would be:

for each pixel in image:
  if pixel[r, g, b] in special:
    return special[pixel[r, g, b]] + pixel[a]
  else:
    return T(pixel)

It's simple but speed has been sub-optimal. I believe there's some way using numpy vectors, but I could not find how.

Important details about the implementation:

I don't care about the original buffer/image (manipulation can be in place)
I can use wxPython, Pillow and NumPy
Order or dimension of the array is not important as long as the buffer keeps the length

The buffer is obtained from a wxPython Bitmap and special and (RG|B)_pal are transformation tables, the end result will become a wxPython Bitmap too. They're obtained like these:

# buffer
bitmap = wx.Bitmap # it's valid wxBitmap here, this is just to let you know it exists
buff = bytearray(bitmap.GetWidth() * bitmap.GetHeight() * 4)
bitmap.CopyToBuffer(buff, wx.BitmapBufferFormat_RGBA)

self.RG_mult= 0.75
self.B_mult = 0.83

self.RG_pal = []
self.B_pal = []

for i in range(0, 256):
    self.RG_pal.append(int(i * self.RG_mult))
    self.B_pal.append(int(i * self.B_mult))

self.special = {
    # RGB: new_RGB
    # Implementation specific for the fastest access
    # with buffer keys are 24bit numbers, with PIL keys are tuples
}

Implementations I tried include direct buffer manipulation:

for x in range(0, bitmap.GetWidth() * bitmap.GetHeight()):
    index = x * 4
    r = buf[index]
    g = buf[index + 1]
    b = buf[index + 2]
    rgb = buf[index:index + 3]
    if rgb in self.special:
        special = self.special[rgb]
        buf[index] = special[0]
        buf[index + 1] = special[1]
        buf[index + 2] = special[2]
    else:
        buf[index] = self.RG_pal[r]
        buf[index + 1] = self.RG_pal[g]
        buf[index + 2] = self.B_pal[b]

Use Pillow with getdata():

pil = Image.frombuffer("RGBA", (bitmap.GetWidth(), bitmap.GetHeight()), buf)
pil_buf = []

for colour in pil.getdata():
    colour_idx = colour[0:3]

    if (colour_idx in self.special):
        special = self.special[colour_idx]
        pil_buf.append((
            special[0],
            special[1],
            special[2],
            colour[3],
        ))
    else:
        pil_buf.append((
            self.RG_pal[colour[0]],
            self.RG_pal[colour[1]],
            self.B_pal[colour[2]],
            colour[3],
        ))

pil.putdata(pil_buf)
buf = pil.tobytes()

Pillow with point() and getdata() (fastest I achieved, more than twice times faster than others)

pil = Image.frombuffer("RGBA", (bitmap.GetWidth(), bitmap.GetHeight()), buf)

r, g, b, a = pil.split()
r = r.point(lambda r: r * self.RG_mult)
g = g.point(lambda g: g * self.RG_mult)
b = b.point(lambda b: b * self.B_mult)
pil = Image.merge("RGBA", (r, g, b, a))

i = 0
for colour in pil.getdata():
    colour_idx = colour[0:3]

    if (colour_idx in self.special):
        special = self.special[colour_idx]
        pil.putpixel(
            (i % bitmap.GetWidth(), i // bitmap.GetWidth()),
            (
                special[0],
                special[1],
                special[2],
                colour[3],
            )
        )
    i += 1

buf = pil.tobytes()

I also tried working with numpy.where but then I could not get it to work. With numpy.apply_along_axis it worked but the performance was terrible. Other tries with numpy I could not access the RGB together, only as separated bands.

user_10639694 · Answer

Pure Numpy Version

This first optimization relies on the fact, that one probably has way less special colors than pixels. I use numpy to do all the inner loops. This works well with images of up to 1MP. If You have multiple images I'd recommend the parallel approach.

Let's define a test case:

import requests
from io import BytesIO
from PIL import Image
import numpy as np

# Load some image, so we have the same
response = requests.get("https://upload.wikimedia.org/wikipedia/commons/4/41/Rick_Astley_Dallas.jpg")
# Make areas of known color
img = Image.open(BytesIO(response.content)).rotate(10, expand=True).rotate(-10,expand=True, fillcolor=(255,255,255)).convert('RGBA')

print("height: %d, width: %d (%.2f MP)"%(img.height, img.width, img.width*img.height/10e6))

height: 5034, width: 5792 (2.92 MP)

Define our special colors

specials = {
    (4,1,6):(255,255,255), 
    (0, 0, 0):(255, 0, 255), 
    (255, 255, 255):(0, 255, 0)
}

Algorithm

def transform_map(img, specials, R_factor, G_factor, B_factor):
    # Your transform
    def transform(x, a):
        a *= x
        return a.clip(0, 255).astype(np.uint8)

    # Convert to array
    img_array = np.asarray(img)
    # Extract channels
    R = img_array.T[0]
    G = img_array.T[1]
    B = img_array.T[2]
    A = img_array.T[3]

    # Find Special colors
    # First, calculate a uniqe hash
    color_hashes = (R + 2**8 * G + 2**16 * B)


    # Find inidices of special colors
    special_idxs = []
    for k, v in specials.items():
        key_arr = np.array(list(k))
        val_arr = np.array(list(v))

        spec_hash = key_arr[0] + 2**8 * key_arr[1] + 2**16 * key_arr[2]
        special_idxs.append(
            {
                'mask': np.where(np.isin(color_hashes, spec_hash)),
                'value': val_arr
            }
        )

    # Apply transform to whole image
    R = transform(R, R_factor)
    G = transform(G, G_factor)
    B = transform(B, B_factor)


    # Replace values where special colors were found
    for idx in special_idxs:
        R[idx['mask']] = idx['value'][0]
        G[idx['mask']] = idx['value'][1]
        B[idx['mask']] = idx['value'][2]

    return Image.fromarray(np.array([R,G,B,A]).T, mode='RGBA')

And finally some bench marks on a Intel Core i5-6300U @ 2.40GHz

import time
times = []
for i in range(10):
    t0 = time.time()
    # Test
    transform_map(img, specials, 1.2, .9, 1.2)
    #
    t1 = time.time()
    times.append(t1-t0)
np.round(times, 2)

print('average run time: %.2f +/-%.2f'%(np.mean(times), np.std(times)))

average run time: 9.72 +/-0.91

EDIT Parallelization

With the same setup as above, we can get a 2x speed increase on large images. (Small ones are faster without numba)

from numba import njit, prange
from numba.core import types
from numba.typed import Dict

# Map dict of special colors or transform over array of pixel values
@njit(parallel=True, locals={'px_hash': types.uint32})
def check_and_transform(img_array, d, T):
    #Save Shape for later
    shape = img_array.shape
    # Flatten image for 1-d iteration
    img_array_flat = img_array.reshape(-1,3).copy()
    N = img_array_flat.shape[0]
    # Replace or map
    for i in prange(N):
        px_hash = np.uint32(0)
        px_hash += img_array_flat[i,0]
        px_hash += types.uint32(2**8) * img_array_flat[i,1] 
        px_hash += types.uint32(2**16) * img_array_flat[i,2]
        
        try:
            img_array_flat[i] = d[px_hash]
        except Exception:
            img_array_flat[i] =  (img_array_flat[i] * T).astype(np.uint8)
    # return image
    return img_array_flat.reshape(shape) 

# Wrapper for function above
def map_or_transform_jit(image: Image, specials: dict, T: np.ndarray):
    # assemble numba typed dict
    d = Dict.empty(
        key_type=types.uint32,
        value_type=types.uint8[:],
    )
    for k, v in specials.items():
        k = types.uint32(k[0] + 2**8 * k[1] + 2**16 * k[2])
        v = np.array(v, dtype=np.uint8)
        d[k] = v
        
    # get rgb channels
    img_arr = np.array(img)
    rgb = img_arr[:,:,:3].copy()
    img_shape = img_arr.shape
    # apply map
    rgb = check_and_transform(rgb, d, T)
    # set color channels
    img_arr[:,:,:3] = rgb
    
    return Image.fromarray(img_arr, mode='RGBA')

# Benchmark
import time
times = []
for i in range(10):
    t0 = time.time()
    # Test
    test_img = map_or_transform_jit(img, specials, np.array([1, .5, .5]))
    #
    t1 = time.time()
    times.append(t1-t0)
np.round(times, 2)

print('average run time: %.2f +/- %.2f'%(np.mean(times), np.std(times)))
test_img

average run time: 3.76 +/- 0.08

Quick pixel manipulation with Pillow and/or NumPy

1 Answers

Pure Numpy Version

EDIT Parallelization