使用 NumPy 将固定调色板应用于图像？

我有一个 RGB 字节的 NumPy 图像，假设它是这个 2x3 图像：

img = np.array([[[  0, 255,   0], [255, 255, 255]],
                [[255,   0, 255], [  0, 255, 255]],
                [[255,   0, 255], [  0,   0,   0]]])

我还有一个调色板，涵盖图像中使用的每种颜色。假设是这个调色板：

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

是否有某种针对调色板索引图像的组合（反之亦然）可以为我提供与此等效的调色板图像？

img_p = np.array([[1, 4],
                  [0, 2],
                  [0, 3]])

作为比较，我知道相反的情况非常简单。palette[img_p]将给出相当于的结果img。我试图找出相反方向是否有类似的方法可以让 NumPy 完成所有繁重的工作。

我知道我可以单独迭代所有图像像素并构建我自己的调色板图像。我希望有一个更优雅的选择。

好的，我实现了下面的各种解决方案，并在一个中等的测试集上运行它们：20 个图像，每个图像 2000x2000 像素，具有 32 元素的三字节颜色调色板。像素被赋予随机调色板索引。所有算法都在相同的图像上运行。

计时结果：

• 大部分为空的查找数组- 0.89秒
• np.search排序方法- 3.20 秒
• Pandas 查找，单个整数- 38.7秒
• 使用 == 然后聚合布尔结果- 66.4 秒
• 将调色板反转为字典并使用 np.apply_along_axis()- 大约 500 秒，基于较小的测试集
• 使用 MultiIndex 进行 Pandas 查找- 大约 3000 秒，基于较小的测试集

鉴于查找数组具有显着的内存损失（如果存在 alpha 通道，则内存损失会令人望而却步），我将采用np.searchsorted方法。查找数组is如果你想在上面使用 RAM，速度会明显加快。

Edit这是一种更快的方法，使用np.searchsorted.

def rev_lookup_by_sort(img, palette):
    M = (1 + palette.max())**np.arange(3)
    p1d, ix = np.unique(palette @ M, return_index=True)
    return ix[np.searchsorted(p1d, img @ M)]

正确性（相当于rev_lookup_by_dict()在下面的原始答案中）：

np.array_equal(
    rev_lookup_by_sort(img, palette),
    rev_lookup_by_dict(img, palette),
)

加速（对于 1000 x 1000 图像和 1000 种颜色的调色板）：

orig = %timeit -o rev_lookup_by_dict(img, palette)
# 2.47 s ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

v2 = %timeit -o rev_lookup_by_sort(img, palette)
# 71.8 ms ± 93.7 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> orig.average / v2.average
34.46

所以这个答案使用np.searchsorted在该尺寸下速度提高了 30 倍。

原答案

最初的镜头给出了一个较慢的版本（希望我们能做得更好）。它使用一个dict，其中键是作为元组的颜色。

def rev_lookup_by_dict(img, palette):
    d = {tuple(v): k for k, v in enumerate(palette)}
    def func(pix):
        return d.get(tuple(pix), -1)
    return np.apply_along_axis(func, -1, img)

img_p = rev_lookup_by_dict(img, palette)

请注意，“未找到颜色”表示为-1 in img_p.

关于您的（修改后的）数据：

>>> img_p
array([[1, 4],
       [0, 2],
       [0, 3]])

更大的例子：

# setup
from math import isqrt

w, h = 1000, 1000
s = isqrt(w * h)
palette = np.random.randint(0, 256, (s, 3))
img = palette[np.random.randint(0, s, (w, h))]

Test:

img_p = rev_lookup_by_dict(img, palette)

>>> np.array_equal(palette[img_p], img)
True

Timing:

%timeit rev_lookup_by_dict(img, palette)
# 2.48 s ± 16.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

这确实很糟糕，但希望我们能做得更好。