Is there a way to optimize this dictionary search?

Question

I have some code which opens an image and prints a representation of it to the terminal using ANSI escape sequences. It takes each pixel in a downsized version of the image and prints a character with a colour that matches it. Not all terminals support RGB output however, so I wanted to implement other colour modes such as 4-bit. I did this with a dictionary lookup table containing the ANSI codes and the RGB values they produce: {(r, g, b) : code, .... I then print the value closest to the pixel colour in terms of euclidean distance.

from PIL import Image, ImageOps
from math import sqrt, inf

def match_colour_to_table(table, colour):
    best_colour = None
    best_diff = inf
    
    for table_colour in table.keys():
        # Calculate the distance between the two colours
        delta = (c - t for c, t in zip(colour, table_colour))
        diff = sqrt(sum((p ** 2 for p in delta)))
        
        if diff < best_diff:
            best_colour = table_colour
            best_diff = diff

    return table[selected_colour]

def term_pixel_4bit(colour):
    colour_table = {
        (0,   0,   0)   : 40,
        (255, 0,   0)   : 41,
        (0,   255, 0)   : 42,
        (255, 255, 0)   : 43,
        (0,   0,   255) : 44,
        (255, 0,   255) : 45,
        (0,   255, 255) : 46,
        (255, 255, 255) : 47,
    }

    code = match_colour_to_table(colour_table, colour)
        
    return f"\033[;{code}m \033[0m"

def term_pixel_256(colour):
    match_colour_to_table(TABLE_256, colour)

    return f"\033[48;5;{code}m ";

def print_image(image, size):
    width, height = size
    image = image.resize(size, resample=1).convert("RGB")
    
    # Print each row of characters
    for y in range(height):            
        row = [term_pixel_256(image.getpixel((x, y)))
               for x in range(width)]
        
        # Reset to avoid trailing colours
        row.append("\033[0m")
        printf("".join(row))

This approach works very well for 4-bit colour, but far less so for 256-colour. I converted the json data at https://jonasjacek.github.io/colors/data.json to a dictionary.

TABLE_256 = {
    (0, 0, 0) : 0, (128, 0, 0) : 1, (0, 128, 0) : 2, 
    ... 
    (228, 228, 228) : 254, (238, 238, 238) : 255
}

It does produce a nice-looking result, but understandably it takes a while to compute. I'm certain there is a much faster way of doing this, but I'm not entirely sure where to begin. Any help would be much appreciated.

Just in case it's needed, here's the calling site:

path  = os.path.join(os.path.dirname(__file__), "")
image = Image.open(path + sys.argv[1])
print_image(image, (100, 50))

Don't know if it would be significant, but you should use [`math.hypot()`](https://docs.python.org/3/library/math.html#math.hypot) to calculate Euclidean distances. Also note that the eye is not equally sensitive to red, green, and blue, so using "distance" in that color space isn't accurate in terms of perception. See [color difference](https://en.wikipedia.org/wiki/Color_difference), — martineau, Dec 14 '20 at 18:14
@martineau Cheers for the response. I switched to one of the algorithms in that wikipedia page. Much better result. I initially did use hypot but in my version of python it only accepts two parameters. — shon-r, Dec 14 '20 at 19:44
Don't know which algorithm you're using now, but noted that in the code in your question, you could leave out the `sqrt()` when calculating `diff` and just use squared distances — it would provide identical results but require less computation. — martineau, Dec 14 '20 at 20:17
I'm using the 'red mean' algorithm, which also allows the `sqrt()` to be left out. I appreciate your help — shon-r, Dec 14 '20 at 20:49
One other thing: It would be faster to use `for table_colour in table:` than `for table_colour in table.keys():` because the latter adds the unnecessary overhead of calling the `keys()` dictionary method. Off-topic: I think the last line of `match_colour_to_table()` should be `return table[best_colour]`. — martineau, Dec 14 '20 at 23:23
Just spotted something else: Instead of calling `getpixel()` on every pixel in the image, you should call [`getdata()`](https://pillow.readthedocs.io/en/latest/reference/Image.html#PIL.Image.Image.getdata) to get them all at once. You can convert the value it returns into rows of pixels the way I do it in [this answer](https://stackoverflow.com/a/40729543/355230). — martineau, Dec 14 '20 at 23:45

score 1 · Answer 1 · answered Dec 14 '20 at 17:43

1

You can try applying functools.lru_cache to term_pixel_4bit

import functools

@functools.lru_cache
def term_pixel_256(colour):
    match_colour_to_table(TABLE_256, colour)

    return f"\033[48;5;{code}m ";

answered Dec 14 '20 at 17:43

This halved the execution time. Cheers! – shon-r Dec 14 '20 at 19:46

score 1 · Accepted Answer · answered Dec 14 '20 at 21:57

You might want to try a KD-Tree to do the nearest neighbor computation in 3-dimensions. Here is an example that employs that with the same color table you referenced. For the small test I have here, it does 50K pixels in about 20 seconds on my machine. Of course, this can be enhanced with lru_cache, which will speed things up to the extent that the exact RGB values are exactly repeated in the sample.

import json
from scipy.spatial import KDTree
from functools import lru_cache
from random import randint

with open('color_table.json', 'r') as f:
    data = json.load(f)

rgbs = [(t['rgb']['r'], t['rgb']['g'], t['rgb']['g']) for t in data]

tree = KDTree(rgbs)  # make the tree

@lru_cache
def match_to_table_tree(tree, colour):
    '''return the index of the colour in the table closest to the colour provided'''
    _, idx = tree.query(colour)
    return idx


test_size=5

test_colors = [(randint(0,255), randint(0,255), randint(0,255)) for t in range(test_size)]

for colour in test_colors:
    idx = match_to_table_tree(tree, colour)
    matched_color = data[idx]['rgb']
    print(f'tested: {colour} -> index {idx} colour: {matched_color}')

Result of 5-color test:

tested: (93, 213, 100) -> index 70 colour: {'r': 95, 'g': 175, 'b': 0}
tested: (37, 204, 3) -> index 22 colour: {'r': 0, 'g': 95, 'b': 0}
tested: (113, 211, 147) -> index 70 colour: {'r': 95, 'g': 175, 'b': 0}
tested: (139, 62, 122) -> index 94 colour: {'r': 135, 'g': 95, 'b': 0}
tested: (35, 106, 107) -> index 22 colour: {'r': 0, 'g': 95, 'b': 0}
[Finished in 0.3s]

This worked very well. Does it in fractions of a second now. Appreciate the help. — shon-r, Dec 15 '20 at 02:01
super. If this answered your question, you could close it out by hitting the check mark to close the issue. — AirSquid, Dec 15 '20 at 02:23

Is there a way to optimize this dictionary search?

2 Answers2

Result of 5-color test: