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

Question

我有一个 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 元素的三字节颜色调色板。像素被赋予运行dom 调色板索引。所有算法都运行处理相同的图像。

计时结果：

- 0.89 秒
- 3.20 秒
- 38.7 秒
- 66.4 秒
- 可能约 500 秒，基于较小的测试集
- 可能约 3000 秒，基于较小的测试集

考虑到查找数组有显着的内存损失（如果有 alpha 通道，则会令人望而却步），我将采用 np.searchsorted 方法。查找数组如果您想在其上花费 RAM，速度会显着加快。

Answer 1

编辑这里有一个使用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)

注意在img_p.

中“找不到颜色”表示为-1

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

>>> 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))]

测试：

img_p = rev_lookup_by_dict(img, palette)

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

时间：

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

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

Answer 2

比字典更快，但具有 64 MB 的查找数组。

d = np.zeros((256,256,256), np.int32)  # 64 MB!
d[tuple(palette.T)] = np.arange(len(palette))
img_p = d[tuple(img.reshape(-1,3).T)].reshape(*img.shape[:2])
# %%timeit 10 loops, best of 5: 25.8 ms per loop (1000 x 1000)

np.testing.assert_equal(img, palette[img_p])

Answer 3

如果除了 NumPy 之外还可以使用 Pandas，则可以使用 Pandas MultiIndex 作为一种稀疏数组：

inverse_palette = pd.Series(np.arange(len(palette)),
                            index=pd.MultiIndex.from_arrays(palette.T)).sort_index()
img_p = np.apply_along_axis(lambda px: inverse_palette[tuple(px)], 2, img)

不过那真的很慢。您可以先将颜色转换为整数来做得更好：

def collapse_bytes(array):
    result = np.zeros(array.shape[:-1], np.uint32)
    for i in range(array.shape[-1]):
        result = result * 256 + array[...,i]
    return result

inverse_palette = pd.Series(np.arange(len(palette)),
                            index=collapse_bytes(palette)).sort_index()
img_p = inverse_palette[collapse_bytes(img).flat].to_numpy()\
                                                 .reshape(img.shape[:-1])

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

Use NumPy to apply a fixed palette to an image?

python

numpy