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

Question

这个问题在this other one之后，而旨在加速以下代码。我构建了一些代码，这些代码从 mxnx 3 获取像素值numpy.ndarray（RGB图像），将像素值与（位置）查找表的像素值进行比较，并输出查找表的像素值 indices ，如下所示：

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

如果您在lut[x]中的多个位置获得img，

错误将无法正确输出。也许更好的策略是将输出格式化为嵌套列表，因为任何可索引的都可以。例如：

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

我只需要获取我在y中查找的像素值的lut索引，然后获取y的{​​{1}}元素以输出值I需要（标量，参见上面的算法摘要）。

Answer 1

你可以使用Kd-tree，这是一个演示：

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