Replace colors in image by closest color in palette using numpy

Question

I have a list of colors, and I have a function closest_color(pixel, colors) where it compares the given pixels' RGB values with my list of colors, and it outputs the closest color from the list.

I need to apply this function to a whole image. When I try to use it pixel by pixel, (by using 2 nested for-loops) it is slow. Is there a better way to achieve this with numpy?

Answer 1

1. Option: Single image evaluation (slow)

Pros

 - any palette any time (flexible)

Cons

 - slow
 - memory for large number of colors in palette
 - not good for batch processing

2. Option: Batch processing (super fast)

Pros
 - super fast (50ms per image), independent of palette size
 - low memory, independent of image size or pallete size
 - ideal for batch processing if palette doesnt change
 - simple code 
Cons
 - requires creation of color cube (once, up to 3 minutes)
 - color cube can contain only one palette

Requirements
 - color cube requires 1.5mb of space on disk in form of compressed np matrix

Option 1:

take image, create pallete object with same size as image, calculate distances, retrieve new image with np.argmin indices

import numpy as np
from PIL import Image
import requests

# get some image
im = Image.open(requests.get("https://upload.wikimedia.org/wikipedia/commons/thumb/7/77/Big_Nature_%28155420955%29.jpeg/800px-Big_Nature_%28155420955%29.jpeg", stream=True).raw)
newsize = (1000, 1000)
im = im.resize(newsize)
# im.show()
im = np.asarray(im)
new_shape = (im.shape[0],im.shape[1],1,3)

# Ignore above
# Now we have image of shape (1000,1000,1,3). 1 is there so its easy to subtract from color container
image = im.reshape(im.shape[0],im.shape[1],1,3)



# test colors
colors = [[0,0,0],[255,255,255],[0,0,255]]

# Create color container 
## It has same dimensions as image (1000,1000,number of colors,3)
colors_container = np.ones(shape=[image.shape[0],image.shape[1],len(colors),3])
for i,color in enumerate(colors):
    colors_container[:,:,i,:] = color



def closest(image,color_container):
    shape = image.shape[:2]
    total_shape = shape[0]*shape[1]

    # calculate distances
    ### shape =  (x,y,number of colors)
    distances = np.sqrt(np.sum((color_container-image)**2,axis=3))

    # get position of the smalles distance
    ## this means we look for color_container position ????-> (x,y,????,3)
    ### before min_index has shape (x,y), now shape = (x*y)
    #### reshaped_container shape = (x*y,number of colors,3)
    min_index = np.argmin(distances,axis=2).reshape(-1)
    # Natural index. Bind pixel position with color_position
    natural_index = np.arange(total_shape)

    # This is due to easy index access
    ## shape is (1000*1000,number of colors, 3)
    reshaped_container = colors_container.reshape(-1,len(colors),3)

    # Pass pixel position with corresponding position of smallest color
    color_view = reshaped_container[natural_index,min_index].reshape(shape[0],shape[1],3)
    return color_view

# NOTE: Dont pass uint8 due to overflow during subtract
result_image = closest(image,colors_container)

Image.fromarray(result_image.astype(np.uint8)).show()

Option 2:

build 256x256x256x3 size color cube based on your palette. In other words, for every existing color assign corresponding palette color that is closest. Save color cube (once/first time). Load color cube. Take image and use every color in image as index in color cube.

import numpy as np
from PIL import Image
import requests
import time
# get some image
im = Image.open(requests.get("https://helpx.adobe.com/content/dam/help/en/photoshop/using/convert-color-image-black-white/jcr_content/main-pars/before_and_after/image-before/Landscape-Color.jpg", stream=True).raw)
newsize = (1000, 1000)
im = im.resize(newsize)
im = np.asarray(im)


### Initialization: Do just once
# Step 1: Define palette
palette = np.array([[255,255,255],[125,0,0],[0,0,125],[0,0,0]])

# Step 2: Create/Load precalculated color cube
try:
    # for all colors (256*256*256) assign color from palette
    precalculated = np.load('view.npz')['color_cube']
except:
    precalculated = np.zeros(shape=[256,256,256,3])
    for i in range(256):
        print('processing',100*i/256)
        for j in range(256):
            for k in range(256):
                index = np.argmin(np.sqrt(np.sum(((palette)-np.array([i,j,k]))**2,axis=1)))
                precalculated[i,j,k] = palette[index]
    np.savez_compressed('view', color_cube = precalculated)
        

# Processing part
#### Step 1: Take precalculated color cube for defined palette and 

def get_view(color_cube,image):
    shape = image.shape[0:2]
    indices = image.reshape(-1,3)
    # pass image colors and retrieve corresponding palette color
    new_image = color_cube[indices[:,0],indices[:,1],indices[:,2]]
   
    return new_image.reshape(shape[0],shape[1],3).astype(np.uint8)

start = time.time()
result = get_view(precalculated,im)
print('Image processing: ',time.time()-start)
Image.fromarray(result).show()