在numpy中向量化非平凡的for循环

时间:2013-10-16 01:31:02

标签: python numpy vectorization

我试图通过矢量化使代码运行得更快,因为我相信python中的循环很慢。我并不完全理解向量化,因此for循环中的切片会给我带来麻烦。

注意:这是针对不允许任何非numpy库的作业。

self.gx = numpy.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
self.gy = numpy.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

def create_vector(self, image):
    """Creates a gradient vector for each pixel in the image

    Returns: vec_data: [mag, angle, avg_value]"""
    vec_data = numpy.zeros((image.shape[0], image.shape[1], 3), dtype='float')

    for y in xrange(1, image.shape[0]-1):
        for x in xrange(1, image.shape[1]-1):

           #Get 3x3 matrix around the pixel
           subimage = image[y-1:y+2,x-1:x+2]

           #Apply sobel operator
           dx = (self.gx*subimage).sum()
           dy = (self.gy*subimage).sum()

           vec_data[y,x,0] = abs(dx) + abs(dy)
           vec_data[y,x,1] = abs(math.atan2(dx,dy))
           vec_data[y,x,2] = subimage.sum()/9 #Average of 3x3 pixels around x, y

    return vec_data

2 个答案:

答案 0 :(得分:4)

使用窗口视图完成代码的直接矢量化:

import numpy as np

image = np.arange(25).reshape(5, 5)
gx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
gm = np.ones((3, 3))/9
rows, cols = image.shape
k_rows, k_cols = gx.shape

from numpy.lib.stride_tricks import as_strided
image_view = as_strided(image, shape=(rows - k_rows + 1, cols - k_cols + 1,
                                      k_rows, k_cols),
                        strides=image.strides*2)
dx = np.einsum('ijkl,kl->ij',image_view, gx)
dy = np.einsum('ijkl,kl->ij',image_view, gy)
dm = np.einsum('ijkl,kl->ij',image_view, gm)

>>> dx
array([[-8, -8, -8],
       [-8, -8, -8],
       [-8, -8, -8]])
>>> dy
array([[-40, -40, -40],
       [-40, -40, -40],
       [-40, -40, -40]])
>>> dm
array([[  6.,   7.,   8.],
       [ 11.,  12.,  13.],
       [ 16.,  17.,  18.]])

从那些你可以构建你想要的输出。

如果遇到性能问题,你的卷积内核是可分离的,即2D卷积可以沿正交轴分成2个1D卷积,这应该比上述解决方案运行得快。

答案 1 :(得分:1)

因此,实际上你通过在移动窗口中进行乘法和求和所做的就是所谓的卷积。 Numpy / Scipy在ndimage模块中具有此功能。因此,您可以获得dxdy的数组,而不是一次只有一个窗口。然后,您可以使用适用于整个magang数组的numpy函数获取avgdxdy图层。因此,分别计算每个层,然后返回三个事物的dstack,如果你想在一个数组中全部:

import numpy as np
from scipy import ndimage

gx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

def create_vector(image):
    """Creates a gradient vector for each pixel in the image

    Returns: vec_data: [mag, angle, avg_value]"""

    #Apply sobel operator using convolution
    dx = ndimage.convolve(gx, image)
    dy = ndimage.convolve(gy, image)

    vec_data_mag = np.abs(dx) + np.abs(dy)
    vec_data_ang = np.abs(np.arctan2(dy, dx))    # are you sure you want abs here?
    vec_data_avg = ndimage.convolve(np.ones(3,3), image)

    return np.dstack([vec_data_mag, vec_data_angl, vec_data_avg])