使用 numpy 获取查找表索引的最快方法

Question

这个问题跟随this other one，并且旨在加速以下代码。我（在帮助下）构建了一些从m x n x 3 @987654323获取像素值的代码@（RGB图像），将像素值与（位置）查找表的像素值进行比较，并输出查找表的像素值索引，如下：

img = np.random.randint(0, 9 , (3, 3, 3))
lut2 = img[0,0:2,:]

output = []
for x in xrange(lut2.shape[0]):
    if lut2[x] in img:
            output.append(np.concatenate(np.where( (img == lut2[x]).sum(axis=2) == 3 )))

print np.array(output)

输出：

[[0 0]
 [0 1]]

bug：如果在img 的多个位置出现lut[x]，将无法正确输出。也许更好的策略是将输出格式化为嵌套列表，因为任何可索引的东西都可以。例如：

[[x_px1, y_px1], [[x_px2a, y_px2a], [x_px2b, y_px2b]]]

其中px2a 和px2b 是img 中具有相同颜色值的独立像素。

我一直在摆弄numpy 索引，但我没有比这更好的了。虽然上述方法部分有效，但由于数组上的迭代，它当然非常慢。如果我给它提供一个普通大小的图像，它会花费不可接受的时间来完成。

有人可以指点我更快的解决方案吗？

---

编辑：我所针对的理想未来算法的总结：

  an image               all pixels replaced          all indices are
  (here, single ----->   by index in lut     ----->   reused in digital_scale
  black pixel)           (here, first pixel)          to output an array of scalars

  [[[0, 0, 0]]]          [[[0, 0]]]                   [[0]]

  array shape:           array shape:                 array shape:

  m X n X 3              m X n X 2                    m X n

最后一个m X n 数组将用于两件事：

• 通过将数组存储在数据库中，以不同的时间间隔绘制数据。
• 对阵列由于一些外部事件引起的变化进行统计分析

注意：基础图像已经是一种绘图，但我无法访问数值数据，所以我需要反转绘图。

---

一些（可选）要求我的代码目标的精度：

查找表通常基于像素值，而不是位置。不过，在这种情况下，位置会更好（我认为）。上面代码中的lut 数组是一个虚拟的线性色标，表示我需要数字化的速度，如下所示：

由于标尺的每一级都是统一的颜色，因此标尺实际上可以缩减为维度数组1 X height X 3（对于RGB、HSV 等）。事实上，我可以删除第一个维度并遍历height X 3 的数组。我用来数字化秤的东西是：

• 它以零为中心
• 它是对称的
• 它是线性的
• 已知最大值和最小值

因此，这个量表的数字表示将是：

digital_scale = np.linspace(-max_speed, max_speed, lut.shape[0])

我只需要得到我在lut中寻找的像素值的y索引，然后是digital_scale的yth元素输出我需要的值（一个标量，见算法摘要以上）。

Answer 1

你可以使用Kd-tree，这里有一个demo：

import numpy as np
from scipy import spatial

H, W = 200, 100

np.random.seed(1)
a = np.random.randint(0, 20, (H, W, 3))
b = np.random.randint(0, 20, (20, 3))
tree = spatial.cKDTree(a.reshape(-1, 3))
res = tree.query_ball_point(b, 0.5, p=1)
print res

输出是：

[[] [577, 17471] [14636, 4515, 13693, 10988, 15013] [16935, 8576, 13286]
 [2443] [7743, 5914] [] [7469, 19736, 13395, 14992, 9083, 15514]
 [1167, 11416] [3903, 4968] [16504, 2996, 10805, 2264] [] [6725]
 [14437, 5888] [17667] [4681, 2545, 6442] [15067, 4533]
 [7876, 2235, 10152, 3288] [15404, 5691, 17216]
 [15586, 9916, 16938, 15931, 4828, 4069]]

检查索引 2 的结果：

rows, cols = np.where(np.all(a == b[None, None, 2], axis=-1))
assert np.all(rows * W + cols == sorted(res[2]))