How do I efficiently update the color of a pixel in a numpy image array if it is closer to one color-CodePudding

Let's say I have a numpy array for an image, img:

import numpy as np

img = np.random.rand(100,100,3) * 255

And a list of colors, palette

white = np.array([255,255,255])
red = np.array([255,0,0])
blue = np.array([0,0,255])
palette = np.array([white, red, blue])

How do I create a new image array, new_img, where every pixel that is closer, through euclidean distance, to white than any of the other colors in the palette is changed to white ([255,255,255]) and every other pixel color is left as is. (It would be nice to be able to change distance functions, but not a hard requirement.)

I am able to do this the naive way with for loops iterating over each pixel, but it's of course a lot slower than it probably needs to be.

CodePudding user response：

First find the distance with every vector and then find the index of the minimum distance.

import numpy as np

img = np.random.rand(100,100,3) * 255
white = np.array([255,255,255])
red = np.array([255,0,0])
blue = np.array([0,0,255])
palette = np.array([white, red, blue])

palette_type = np.linalg.norm(img[...,None,:]-palette[:,:], axis=-1).argmin(axis=-1) ## 0 stands for white, 1 for red and 2 for blue.

new_img  = palette[palette_type]

CodePudding user response：

Thanks to Murali for the solution that immensely helped. I took that:

import numpy as np

img = np.random.rand(100,100,3) * 255
white = np.array([255,255,255])
red = np.array([255,0,0])
blue = np.array([0,0,255])
palette = np.array([white, red, blue])

palette_type = np.linalg.norm(img[...,None,:]-palette[:,:], axis=-1).argmin(axis=-1) ## 0 stands for white, 1 for red and 2 for blue.

new_img  = palette[palette_type]

And added a few lines to take it a little further by making the white pixels transparent and kept all other pixels as their original color:

mask = np.all(new_img == white, axis=-1)
img_copy = img.copy()
img_copy[mask] = white
h,w, _ = img_copy.shape
transparent_img = np.dstack((img_copy, np.zeros((h,w),dtype=np.uint8) 255))
masked_white = (transparent_img[:, :, 0:3] == white).all(2)
transparent_img[masked_white] = [0,0,0,0]