CODEWITHSUNDEEP
pythonnumpyimage-processingpython-imaging-libraryquantization
Fawfulcopter
Fawfulcopter

Reputation: 423

Faster implementation to quantize an image with an existing palette?

I am using Python 3.6 to perform basic image manipulation through Pillow. Currently, I am attempting to take 32-bit PNG images (RGBA) of arbitrary color compositions and sizes and quantize them to a known palette of 16 colors. Optimally, this quantization method should be able to leave fully transparent (A = 0) pixels alone, while forcing all semi-transparent pixels to be fully opaque (A = 255). I have already devised working code that performs this, but I wonder if it may be inefficient:

import math
from PIL import Image

# a list of 16 RGBA tuples
palette = [
    (0, 0, 0, 255),
    # ...
    ]

with Image.open('some_image.png').convert('RGBA') as img:
    for py in range(img.height):
        for px in range(img.width):
            pix = img.getpixel((px, py))

            if pix[3] == 0:  # Ignore fully transparent pixels
                continue

            # Perform exhaustive search for closest Euclidean distance
            dist = 450
            best_fit = (0, 0, 0, 0)
            for c in palette:
                if pix[:3] == c:  # If pixel matches exactly, break
                    best_fit = c
                    break
                tmp = sqrt(pow(pix[0]-c[0], 2) + pow(pix[1]-c[1], 2) + pow(pix[2]-c[2], 2))
                if tmp < dist:
                    dist = tmp
                    best_fit = c
            img.putpixel((px, py), best_fit + (255,))
    img.save('quantized.png')

I think of two main inefficiencies of this code:

Is there a faster method to do this?

I've noted that Pillow has a native function Image.quantize() that seems to do exactly what I want. But as it is coded, it forces dithering in the result, which I do not want. This has been brought up in another StackOverflow question. The answer to that question was simply to extract the internal Pillow code and tweak the control variable for dithering, which I tested, but I find that Pillow corrupts the palette I give it and consistently yields an image where the quantized colors are considerably darker than they should be.

Image.point() is a tantalizing method, but it only works on each color channel individually, where color quantization requires working with all channels as a set. It'd be nice to be able to force all of the channels into a single channel of 32-bit integer values, which seems to be what the ill-documented mode "I" would do, but if I run img.convert('I'), I get a completely greyscale result, destroying all color.

An alternative method seems to be using NumPy and altering the image directly. I've attempted to create a lookup table of RGB values, but the three-dimensional indexing of NumPy's syntax is driving me insane. Ideally I'd like some kind of code that works like this:

img_arr = numpy.array(img)

# Find all unique colors
unique_colors = numpy.unique(arr, axis=0)

# Generate lookup table
colormap = numpy.empty(unique_colors.shape)
for i, c in enumerate(unique_colors):
    dist = 450
    best_fit = None
    for pc in palette:
        tmp = sqrt(pow(c[0] - pc[0], 2) + pow(c[1] - pc[1], 2) + pow(c[2] - pc[2], 2))
        if tmp < dist:
            dist = tmp
            best_fit = pc
    colormap[i] = best_fit

# Hypothetical pseudocode I can't seem to write out
for iy in range(arr.size):
for ix in range(arr[0].size):
    if arr[iy, ix, 3] == 0: # Skip transparent
        continue
    index = # Find index of matching color in unique_colors, somehow
    arr[iy, ix] = colormap[index]

I note with this hypothetical example that numpy.unique() is another slow operation, since it sorts the output. Since I cannot seem to finish the code the way I want, I haven't been able to test if this method is faster anyway.

I've also considered attempting to flatten the RGBA axis by converting the values to a 32-bit integer and desiring to create a one-dimensional lookup table with the simpler index:

def shift(a):
    return a[0] << 24 | a[1] << 16 | a[2] << 8 | a[3]

img_arr = numpy.apply_along_axis(shift, 1, img_arr)

But this operation seemed noticeably slow on its own.

I would prefer answers that involve only Pillow and/or NumPy, please. Unless using another library demonstrates a dramatic computational speed increase over any PIL- or NumPy-native solution, I don't want to import extraneous libraries to do something these two libraries should be reasonably capable of on their own.

Upvotes: 1

Views: 4222

Answers (1)

Ole Tange
Ole Tange

Reputation: 33740

for loops should be avoided for speed.

I think you should make a tensor like:

d2[x,y,color_index,rgb] = distance_squared

where rgb = 0..2 (0 = r, 1 = g, 2 = b).

Then compute the distance:

d[x,y,color_index] =
  sqrt(sum(rgb,d2))

Then select the color_index with the minimal distance:

c[x,y] = min_index(color_index, d)

Finally replace alpha as needed:

alpha = ceil(orig_image.alpha)
img = c,alpha

Upvotes: 1

Related Questions