为我的代码矢量化ndimage函数

时间:2013-10-04 13:28:52

标签: python python-2.7 numpy scipy

我希望能够对此代码进行矢量化:

def sobHypot(rec):
    a, b, c = rec.shape
    hype = np.ones((a,b,c))

    for i in xrange(c):
        x=ndimage.sobel(abs(rec[...,i])**2,axis=0, mode='constant')
        y=ndimage.sobel(abs(rec[...,i])**2,axis=1, mode='constant')
        hype[...,i] = np.hypot(x,y)
        hype[...,i] = hype[...,i].mean()

    index = hype.argmax()
    return index

其中rec,shape返回(1024,1024,20)

1 个答案:

答案 0 :(得分:0)

以下是使用sobel滤镜避免for循环的方法:

import numpy as np
from scipy.ndimage import sobel

def sobHypot_vec(rec):
    r = np.abs(rec)
    x = sobel(r, 0, mode='constant')
    y = sobel(r, 1, mode='constant')
    h = np.hypot(x, y)
    h = np.apply_over_axes(np.mean, h, [0,1])
    return h.argmax()

我不确定sobel滤镜在您的应用中是否特别必要,如果没有您特定的20层“图像”,这很难测试,但您可以尝试使用np.gradient而不是运行sobel两次。优点是gradient在三个维度上运行。您可以忽略第三个中的组件,并采用前两个中的组件。这似乎很浪费,但在我的测试中实际上仍然更快。

对于各种随机生成的图片r = np.random.rand(1024,1024,20) + np.random.rand(1024,1024,20)*1j,这给出了与您的代码相同的答案,但要测试它以确保,并且可能摆弄np.gradientdx, dy个参数

def grad_max(rec):
    g = np.gradient(np.abs(rec))[:2]  # ignore derivative in third dimension
    h = np.hypot(*g)
    h = np.apply_over_axes(np.mean, h, [0,1]) # mean along first and second dimension 
    return h.argmax()

使用此代码进行计时:

def sobHypot_clean(rec):
    rs = rec.shape
    hype = np.ones(rs)
    r = np.abs(rec)
    for i in xrange(rs[-1]):
        ri = r[...,i]
        x = sobel(ri, 0, mode='constant')
        y = sobel(ri, 1, mode='constant')
        hype[...,i] = np.hypot(x,y).mean()
    return hype.argmax()

定时:

In [1]: r = np.random.rand(1024,1024,20) + np.random.rand(1024,1024,20)*1j

# Original Post
In [2]: timeit sobHypot(r)
1 loops, best of 3: 9.85 s per loop

#cleaned up a bit:
In [3]: timeit sobHypot_clean(r)
1 loops, best of 3: 7.64 s per loop

# vectorized:
In [4]: timeit sobHypot_vec(r)
1 loops, best of 3: 5.98 s per loop

# using np.gradient:
In [5]: timeit grad_max(r)
1 loops, best of 3: 4.12 s per loop

请在您自己的图像上测试任何这些功能,以确保它们提供所需的输出,因为不同类型的阵列可能与我所做的简单随机测试的反应不同。