Replacing RGB values in numpy array by integer is extremely slow

Question

I want to replace the rgb values of a numpy array to single integer representations. My code works but it's too slow, I am iterating over every element right now. Can I speed this up? I am new to numpy.

from skimage import io

# dictionary of color codes for my rgb values
_color_codes = {
    (255, 200, 100): 1,
    (223, 219, 212): 2,
    ...
}

# get the corresponding color code for the rgb vector supplied
def replace_rgb_val(rgb_v):
    rgb_triple = (rgb_v[0], rgb_v[1], rgb_v[2])
    if rgb_triple in _color_codes:
        return _color_codes[rgb_triple]
    else:
        return -1

# function to replace, this is where I iterate
def img_array_to_single_val(arr):
    return np.array([[replace_rgb_val(arr[i][j]) for j in range(arr.shape[1])] for i in range(arr.shape[0])])


# my images are square so the shape of the array is (n,n,3)
# I want to change the arrays to (n,n,1)
img_arr = io.imread(filename)
# this takes from ~5-10 seconds, too slow!
result = img_array_to_single_val(img_arr)

Answer 1

Replace the color values the other way round. Look for each RGB-triple, and set the corresponding index in a new array:

def img_array_to_single_val(arr, color_codes):
    result = numpy.ndarray(shape=arr.shape[:2], dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        result[(arr==rgb).all(2)] = idx
    return result

Let's take the color-index assignment apart: First arr==rgb compares each pixel-rgb-values with the list rgb, leading to a n x n x 3 - boolean array. Only if all three color-parts are the same, we found a match, so .all(2) reduces the last axis, leeding to a n x n - boolean array, with True for every pixel matching rgb. Last step is, to use this mask to set the index of the corresponding pixels.

Even faster, it might be, to first convert the RGB-array to int32, and then do the index translation:

def img_array_to_single_val(image, color_codes):
    image = image.dot(numpy.array([65536, 256, 1], dtype='int32'))
    result = numpy.ndarray(shape=image.shape, dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        rgb = rgb[0] * 65536 + rgb[1] * 256 + rgb[2]
        result[arr==rgb] = idx
    return result

For really large or many images you should first create a direct color mapping:

color_map = numpy.ndarray(shape=(256*256*256), dtype='int32')
color_map[:] = -1
for rgb, idx in color_codes.items():
    rgb = rgb[0] * 65536 + rgb[1] * 256 + rgb[2]
    color_map[rgb] = idx

def img_array_to_single_val(image, color_map):
    image = image.dot(numpy.array([65536, 256, 1], dtype='int32'))
    return color_map[image]

Answer 2

Two fully vectorized solutions could be suggested here.

Approach #1: Using NumPy's powerful broadcasting capability -

# Extract color codes and their IDs from input dict
colors = np.array(_color_codes.keys())
color_ids = np.array(_color_codes.values())

# Initialize output array
result = np.empty((img_arr.shape[0],img_arr.shape[1]),dtype=int)
result[:] = -1

# Finally get the matches and accordingly set result locations
# to their respective color IDs
R,C,D = np.where((img_arr == colors[:,None,None,:]).all(3))
result[C,D] = color_ids[R]

Approach #2: Using cdist from scipy.spatial.distance one can replace the final steps from approach #1, like so -

from scipy.spatial.distance import cdist

R,C = np.where(cdist(img_arr.reshape(-1,3),colors)==0)
result.ravel()[R] = color_ids[C]

Answer 3

Going through every pixel manually and creating a dict of 256**3 items just to get another color palette seems strange to me if you are not after a specific effect you want to create. If you just want to flatten the image to integer values you can use the skimage rg2gray(img) function. This will give you the pictures luminance.

You can use pylabs colormaps to get another representation:

import matplotlib.pylab as plt
import skimage 
import matplotlib.cm as cm

img = io.imread("Fox.jpg")
gray_img = skimage.color.rgb2gray(img)
plt.imshow(img, cmap=cm.Jet)
plt.show()

Answer 4

The last solution in the accepted answer by @Daniel is fantastic for large images. And it's possible to make it even faster by reshaping the input image to a 2D array and performing a single vectorized calculation to convert the RGB values to a single integer value using only element-wise multiplication and addition:

def compute_colormap(colors: List[Tuple[int, int, int]]) -> np.ndarray:
    colormap = np.zeros((256**3,), dtype=np.int32)
    for i, color in enumerate(colors):
        colormap[color[0] * 256**2 + color[1] * 256 + color[2]] = i
    return colormap

def image2label(image: np.ndarray, colormap: np.ndarray):
    flat = image.reshape(-1, 3)
    label = colormap[flat[:, 0] * 256**2 + flat[:, 1] * 256 + flat[:, 2]]
    return label.reshape(image.shape[:-1])

Benchmarks for a large (1024x14640x3) image:

Without flattening and using dot-product (original accepted answer): 170 ms ± 541 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

This answer: 57.1 ms ± 309 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

If you are storing these labels to disk rather than converting during runtime — which may become a bottleneck in your data loaders — you can change the type of the labels or colormap to uint8 or int16 (depending on how many classes you have) which will save disk space and greatly reduce load times.