Replace colors in image by closest color in palette using numpy

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?


  • 1. Option: Single image evaluation (slow)

     - any palette any time (flexible)
     - slow
     - memory for large number of colors in palette
     - not good for batch processing

    2. Option: Batch processing (super fast)

     - 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 
     - requires creation of color cube (once, up to 3 minutes)
     - color cube can contain only one palette
     - 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 ="", stream=True).raw)
    newsize = (1000, 1000)
    im = im.resize(newsize)
    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)

    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 ="", 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
        # for all colors (256*256*256) assign color from palette
        precalculated = np.load('view.npz')['color_cube']
        precalculated = np.zeros(shape=[256,256,256,3])
        for i in range(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)