有没有办法优化这个字典搜索？

Question

我有一些代码可以打开图像并使用 ANSI 转义序列将其表示形式打印到终端。它采用缩小版图像中的每个像素，并打印出颜色与之匹配的字符。然而，并非所有终端都支持 RGB 输出，因此我想实现其他颜色模式，例如 4 位。我通过包含 ANSI 代码和它们产生的 RGB 值的字典查找 table 来做到这一点：{(r, g, b) : code, ...。然后我根据欧氏距离打印最接近像素颜色的值。

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"3[;{code}m 3[0m"

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

    return f"3[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("3[0m")
        printf("".join(row))

这种方法对于 4 位色效果很好，但对于 256 色效果就差很多了。我将 https://jonasjacek.github.io/colors/data.json 处的 json 数据转换为字典。

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

它确实产生了一个漂亮的结果，但可以理解它需要一段时间来计算。我确信有一种更快的方法可以做到这一点，但我不完全确定从哪里开始。任何帮助将不胜感激。

为了以防万一，这里是调用站点：

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

Answer 1

您可以尝试将 functools.lru_cache 应用到 term_pixel_4bit

import functools

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

    return f"3[48;5;{code}m ";

Answer 2

您可能想尝试使用 KD 树在 3 维中进行最近邻计算。这是一个使用与您引用的相同颜色 table 的示例。对于我这里的小测试，它在我的机器上大约 20 秒内完成 50K 像素。当然，这可以通过 lru_cache 来增强，这将加快速度，使样本中的 RGB 值完全重复。

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}')

5色测试结果：

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]

有没有办法优化这个字典搜索？

Is there a way to optimize this dictionary search?

python

lookup

terminal

optimization

python-imaging-library

5色测试结果：