编辑我保留了下面我面临的更复杂的问题,但我对np.take
的问题可以更好地总结如下。假设您有一个形状为img
的数组(planes, rows)
,另一个形状为lut
的数组(planes, 256)
,并且您希望使用它们来创建一个新数组out
形状(planes, rows)
,其中out[p,j] = lut[p, img[p, j]]
。这可以通过如下的花式索引来实现:
In [4]: %timeit lut[np.arange(planes).reshape(-1, 1), img]
1000 loops, best of 3: 471 us per loop
但是,如果你使用take和python循环来代替花哨的索引而不是planes
那么事情可以大大加快:
In [6]: %timeit for _ in (lut[j].take(img[j]) for j in xrange(planes)) : pass
10000 loops, best of 3: 59 us per loop
lut
和img
可以重新排列,以便在没有python循环的情况下完成整个操作,但使用numpy.take
(或替代方法)而不是传统的花式索引保持速度优势?
原始问题
我有一组想要在图像上使用的查找表(LUT)。保持LUT的数组的形状为(planes, 256, n)
,图像的形状为(planes, rows, cols)
。两者都是dtype = 'uint8'
,匹配LUT的256
轴。我们的想法是通过LUT的p
平面中的每个n
LUT运行图像的p
平面。
如果我的lut
和img
如下:
planes, rows, cols, n = 3, 4000, 4000, 4
lut = np.random.randint(-2**31, 2**31 - 1,
size=(planes * 256 * n // 4,)).view('uint8')
lut = lut.reshape(planes, 256, n)
img = np.random.randint(-2**31, 2**31 - 1,
size=(planes * rows * cols // 4,)).view('uint8')
img = img.reshape(planes, rows, cols)
在使用像这样的花哨索引之后,我可以实现我的目标
out = lut[np.arange(planes).reshape(-1, 1, 1), img]
它给了我一个形状(planes, rows, cols, n)
的数组,其中out[i, :, :, j]
保持i
的{{1}}平面贯穿img
- LUT LUT的j
- 平面......
一切都很好,除此之外:
i
这是完全不可接受的,特别是因为我使用In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img]
1 loops, best of 3: 5.65 s per loop
以下所有不太好看的替代方案比运行得更快:
单个平面上的单个LUT快速运行x70:
np.take
运行所有所需组合的python循环几乎完成了x6:
In [2]: %timeit np.take(lut[0, :, 0], img[0])
10 loops, best of 3: 78.5 ms per loop
即使在LUT和图像中运行所有平面组合,然后丢弃In [2]: %timeit for _ in (np.take(lut[j, :, k], img[j]) for j in xrange(planes) for k in xrange(n)) : pass
1 loops, best of 3: 947 ms per loop
不需要的平面,也比花哨索引更快:
planes**2 - planes
我能够提出的最快组合有一个迭代飞机的python循环并更快地完成x13:
In [2]: %timeit np.take(lut, img, axis=1)[np.arange(planes), np.arange(planes)]
1 loops, best of 3: 3.79 s per loop
当然,问题是如果没有任何python循环,In [2]: %timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
1 loops, best of 3: 434 ms per loop
没有办法做到这一点?理想情况下,无论需要重塑还是调整大小都应该在LUT上进行,而不是图像,但我愿意接受人们可以提出的任何内容...
答案 0 :(得分:5)
我要说的全部拳头我真的很喜欢你的问题。如果不重新排列LUT
或IMG
,则以下解决方案有效:
%timeit a=np.take(lut, img, axis=1)
# 1 loops, best of 3: 1.93s per loop
但是从结果中你必须查询对角线:a [0,0],a [1,1],a [2,2];得到你想要的。我试图找到一种方法只为对角线元素做索引,但仍然无法管理。
以下是重新排列LUT
和IMG
的一些方法:
如果IMG
中的索引为0-255,第一个平面,第二个平面为256-511,第三个平面为512-767,则以下情况有效,但这会阻止您使用{{1}这可能是个大问题......:
'uint8'
在我的机器中,你的解决方案仍然是最好的选择,而且非常合适,因为你只需要对角线评估,即plane1-plane1,plane2-plane2和plane3-plane3:
lut2 = lut.reshape(-1,4)
%timeit np.take(lut2,img,axis=0)
# 1 loops, best of 3: 716 ms per loop
# or
%timeit np.take(lut2, img.flatten(), axis=0).reshape(3,4000,4000,4)
# 1 loops, best of 3: 709 ms per loop
我希望这可以为您提供更好的解决方案。最好使用%timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass
# 1 loops, best of 3: 677 ms per loop
和类似flatten()
或np.apply_over_axes()
的方法寻找更多选项,这似乎很有希望。
我使用下面的代码生成数据:
np.apply_along_axis()