Calculating the nearest neighbour in a 2d grid using multilevel solution

Viewed 228

I have a problem where in a grid of x*y size I am provided a single dot, and I need to find the nearest neighbour. In practice, I am trying to find the closest dot to the cursor in pygame that crosses a color distance threshold that is calculated as following:

sqrt(((rgb1[0]-rgb2[0])**2)+((rgb1[1]-rgb2[1])**2)+((rgb1[2]-rgb2[2])**2))

So far I have a function that calculates the different resolutions for the grid and reduces it by a factor of two while always maintaining the darkest pixel. It looks as following:

from PIL import Image
from typing import Dict
import numpy as np

#we input a pillow image object and retrieve a dictionary with every grid version of the 3 dimensional array:
def calculate_resolutions(image: Image) -> Dict[int, np.ndarray]: 
    resolutions = {}
    #we start with the highest resolution image, the size of which we initially divide by 1, then 2, then 4 etc.:
    divisor = 1
    #reduce the grid by 5 iterations
    resolution_iterations = 5
    for i in range(resolution_iterations):
        pixel_lookup = image.load() #convert image to PixelValues object, which allows for pixellookup via [x,y] index
        #calculate the resolution of the new grid, round upwards:
        resolution = (int((image.size[0] - 1) // divisor + 1), int((image.size[1] - 1) // divisor + 1))
        #generate 3d array with new grid resolution, fill in values that are darker than white:
        new_grid = np.full((resolution[0],resolution[1],3),np.array([255,255,255]))
        for x in range(image.size[0]):
            for y in range(image.size[1]):
                if not x%divisor and not y%divisor:
                    darkest_pixel = (255,255,255)
                    x_range = divisor if x+divisor<image.size[0] else (0 if image.size[0]-x<0 else image.size[0]-x)
                    y_range = divisor if y+divisor<image.size[1] else (0 if image.size[1]-y<0 else image.size[1]-y)
                    for x_ in range(x,x+x_range):
                        for y_ in range(y,y+y_range):
                            if pixel_lookup[x_,y_][0]+pixel_lookup[x_,y_][1]+pixel_lookup[x_,y_][2] < darkest_pixel[0]+darkest_pixel[1]+darkest_pixel[2]:
                                darkest_pixel = pixel_lookup[x_,y_]
                    if darkest_pixel != (255,255,255):
                        new_grid[int(x/divisor)][int(y/divisor)] = np.array(darkest_pixel)
        resolutions[i] = new_grid
        divisor = divisor*2
    return resolutions

This is the most performance efficient solution I was able to come up with. If this function is run on a grid that continually changes, like a video with x fps, it will be very performance intensive. I also considered using a kd-tree algorithm that simply adds and removes any dots that happen to change on the grid, but when it comes to finding individual nearest neighbours on a static grid this solution has the potential to be more resource efficient. I am open to any kinds of suggestions in terms of how this function could be improved in terms of performance.

Now, I am in a position where for example, I try to find the nearest neighbour of the current cursor position in a 100x100 grid. The resulting reduced grids are 50^2, 25^2, 13^2, and 7^2. In a situation where a part of the grid looks as following:

example grid

And I am on the aggregation step where a part of the grid consisting of six large squares, the black one being the current cursor position and the orange dots being dots where the color distance threshold is crossed, I would not know which diagonally located closest neighbour I would want to pick to search next. In this case, going one aggregation step down shows that the lower left would be the right choice. Depending on how many grid layers I have this could result in a very large error in terms of the nearest neighbour search. Is there a good way how I can solve this problem? If there are multiple squares that show they have a relevant location, do I have to search them all in the next step to be sure? And if that is the case, the further away I get the more I would need to make use of math functions such as the pythagorean theorem to assert whether the two positive squares I find are overlapping in terms of distance and could potentially contain the closest neighbour, which would start to be performance intensive again if the function is called frequently. Would it still make sense to pursue this solution over a regular kd tree? For now the grid size is still fairly small (~800-600) but if the grid gets larger the performance may start suffering again. Is there a good scalable solution to this problem that could be applied here?

0 Answers
Related