[SOLVED] Finding nearest pixel in defined color space – quick implementation using numpy

Issue

I have been working on a task, where I implemented median cut for image quantization – representing the whole image by only limited set of pixels. I implemented the algorithm and now I am trying to implement the part, where I assign each pixel to a representant from the set found by median cut. So, I have variable ‘color_space’, which is 2d ndarray of shape (n,3), where n is the number of representatives. Then I have variable ‘img’, which is the original image of shape (rows, columns, 3).

Now I want to find the nearest pixel (bin) for each pixel from the image based on euclidean distance. I was able to come with this solution:

    for row in range(img.shape[0]):
       for column in range(img.shape[1]):
          img[row][column] = color_space[np.linalg.norm(color_space - img[row][column], axis=1).argmin()]

What it does is, that for each pixel from the image, it computes the vector if distances from each of the bins and then it takes the closest one.
Problem is, that this solution is quite slow and I would like to vectorize it – instead of getting vector for each pixel, I would like to get a matrix, where for example first row would be the first vector of distances computed in my code etc…

This problem could be converted into a problem, where I want to do a matrix multiplication, but instead of getting dot product of two vectors, I would get their euclidean distance. Is there some good approach to such problems? Some general solution in numpy, if we want to do ‘matrix multiplication’ in numpy, but the function Rn x Rn -> R does not need to be dot product, but for example euclidean distance. Of course, for the multiplication, the original image should be resized to (row*columns, 3), but that is a detail.

I have been studying the documentation and searching internet, but didn’t find any good approach.

Please note that I don’t want others to solve my assignment, the solution I came up with is totally ok, I am just curious whether I could speed it up as I try to learn numpy properly.

Thanks for any advices!

Solution

Below is MWE for vectorizing your problem. See comments for explanation.

import numpy

# these are just random array declaration to work with. 
image = numpy.random.rand(32, 32, 3)
color_space = numpy.random.rand(10,3)

# your code. I modified it to pick indexes
result = numpy.zeros((32,32))
for row in range(image.shape[0]):
    for column in range(image.shape[1]):
        result[row][column] = numpy.linalg.norm(color_space - image[row][column], axis=1).argmin()
result = result.astype(numpy.int)

# here we reshape for broadcasting correctly. 
image = image.reshape(1,32,32,3)
color_space = color_space.reshape(10, 1,1,3)

# compute the norm on last axis, which is RGB values
result_norm = numpy.linalg.norm(image-color_space, axis=3)

# now compute the vectorized argmin 
result_vectorized = result_norm.argmin(axis=0)

print(numpy.allclose(result, result_vectorized))

Eventually, you can get the correct solution by doing color_space[result]. You may have to remove the extra dimensions that you add in color space to get correct shapes in this final operation.

Answered By – Umang Gupta

Answer Checked By – Cary Denson (BugsFixing Admin)

Leave a Reply

Your email address will not be published. Required fields are marked *