刚性转换 - Python - 加速

Rigid transformation - Python - speedup

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

我需要为给定图像中的每个 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 倍)。

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

您可以使用 np.indicesnp.rollaxis 生成 3D 数组,其中 coords[i, j] == [i, j]。这里坐标需要切换

然后你所做的就是附加你要求的1,然后使用@

coords_ext = np.empty((Y, X, 3))
coords_ext[...,[1,0]] = np.rollaxis(np.indices((Y, X)), 0, start=3)
coords_ext[...,2] = 1

# convert to column vectors and back for matmul broadcasting
result = (matrix @ coords_ext[...,None])[...,0]

# or alternatively, work with row vectors and do it in the other order
result = coords_ext @ matrix.T

借用 , the idea of creating 1D arrays instead of the 2D or 3D meshes and using them with on-the-fly broadcasted 操作来提高内存效率,从而获得性能优势,这是一种方法 -

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

如果您使用 xsys 覆盖 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