python - 刚性转换 - Python - 加速

Question

我有以下关于计算刚性变换的更快方法的问题(是的，我知道我可以简单地使用库，但需要自己编写代码)。

我需要为给定图像中的每个 x,y 计算 x' 和 y'。我的主要瓶颈是所有坐标的点积(之后的插值不是问题)。目前我实现了三个选项:

1. 列表理解

   result = np.array([[np.dot(matrix, np.array([x, y, 1])) for x in xs] for y in ys])
   

2. 简单的双for 循环

   result = np.zeros((Y, X, 3))
   for x in xs:
       for y in ys:
           result[y, x, :] = np.dot(matrix, np.array([x, y, 1]))
   

3. np.ndenumerate

   result = np.zeros((Y, X, 3))
   for (y, x), value in np.ndenumerate(image):
       result[y, x, :] = np.dot(matrix, np.array([x, y, 1]))
   

在 512x512 图像中最快的方法是列表理解(比 np.ndenumerate 快 1.5 倍，比 double for 循环快 1.1 倍)。

有什么方法可以更快地做到这一点？

Answer 1

借自this post ，创建 1D 数组而不是 2D 或 3D 网格并将它们用于即时 broadcasted 的想法操作以提高内存效率，从而获得性能优势，这是一种方法 -

out = ys[:,None,None]*matrix[:,1] + xs[:,None]*matrix[:,0] + matrix[:,2]

如果您使用 xs 和 ys 覆盖 512x512 大小的图像的所有索引，我们将使用 np 创建它们.arange，像这样 -

ys = np.arange(512)
xs = np.arange(512)

运行时测试

函数定义-

def original_listcomp_app(matrix, X, Y): # Original soln with list-compr. 
    ys = np.arange(Y)
    xs = np.arange(X)
    out = np.array([[np.dot(matrix, np.array([x, y, 1])) for x in xs] \
                                                           for y in ys])
    return out    

def indices_app(matrix, X, Y):        # @Eric's soln  
    coords_ext = np.empty((Y, X, 3))
    coords_ext[...,[1,0]] = np.rollaxis(np.indices((Y, X)), 0, start=3)
    coords_ext[...,2] = 1    
    result = np.matmul(coords_ext,matrix.T)
    return result

def broadcasting_app(matrix, X, Y):  # Broadcasting based
    ys = np.arange(Y)
    xs = np.arange(X)
    out = ys[:,None,None]*matrix[:,1] + xs[:,None]*matrix[:,0] + matrix[:,2]
    return out

时间和验证-

In [242]: # Inputs
     ...: matrix = np.random.rand(3,3)
     ...: X,Y = 512, 512
     ...: 

In [243]: out0 = original_listcomp_app(matrix, X, Y)
     ...: out1 = indices_app(matrix, X, Y)
     ...: out2 = broadcasting_app(matrix, X, Y)
     ...: 

In [244]: np.allclose(out0, out1)
Out[244]: True

In [245]: np.allclose(out0, out2)
Out[245]: True

In [253]: %timeit original_listcomp_app(matrix, X, Y)
1 loops, best of 3: 1.32 s per loop

In [254]: %timeit indices_app(matrix, X, Y)
100 loops, best of 3: 16.1 ms per loop

In [255]: %timeit broadcasting_app(matrix, X, Y)
100 loops, best of 3: 9.64 ms per loop