你可以用as_strided函数与numpy的广播相结合来做一些惊人的事情。以下是您的函数的两个版本:import numpy as np
from numpy.lib.stride_tricks import as_strided
def sumsqdiff(input_image, template, valid_mask=None):
if valid_mask is None:
valid_mask = np.ones_like(template)
total_weight = valid_mask.sum()
window_size = template.shape
ssd = np.empty((input_image.shape[0] - window_size[0] + 1,
input_image.shape[1] - window_size[1] + 1))
for i in xrange(ssd.shape[0]):
for j in xrange(ssd.shape[1]):
sample = input_image[i:i + window_size[0], j:j + window_size[1]]
dist = (template - sample) ** 2
ssd[i, j] = (dist * valid_mask).sum()
return ssd
def sumsqdiff2(input_image, template, valid_mask=None):
if valid_mask is None:
valid_mask = np.ones_like(template)
total_weight = valid_mask.sum()
window_size = template.shape
# Create a 4-D array y, such that y[i,j,:,:] is the 2-D window
# input_image[i:i+window_size[0], j:j+window_size[1]]
y = as_strided(input_image,
shape=(input_image.shape[0] - window_size[0] + 1,
input_image.shape[1] - window_size[1] + 1,) +
window_size,
strides=input_image.strides * 2)
# Compute the sum of squared differences using broadcasting.
ssd = ((y - template) ** 2 * valid_mask).sum(axis=-1).sum(axis=-1)
return ssd
这里有一个ipython会话来比较它们。在
我将用于演示的模板:
^{pr2}$
我们可以检查结果的小输入:In [73]: x
Out[73]:
array([[ 0., 1., 2., 3., 4., 5., 6.],
[ 7., 8., 9., 10., 11., 12., 13.],
[ 14., 15., 16., 17., 18., 19., 20.],
[ 21., 22., 23., 24., 25., 26., 27.],
[ 28., 29., 30., 31., 32., 33., 34.]])
将这两个函数应用于x,并检查是否得到相同的结果:In [74]: sumsqdiff(x, template)
Out[74]:
array([[ 856., 1005., 1172., 1357., 1560.],
[ 2277., 2552., 2845., 3156., 3485.],
[ 4580., 4981., 5400., 5837., 6292.]])
In [75]: sumsqdiff2(x, template)
Out[75]:
array([[ 856., 1005., 1172., 1357., 1560.],
[ 2277., 2552., 2845., 3156., 3485.],
[ 4580., 4981., 5400., 5837., 6292.]])
现在输入更大的“图像”:In [76]: z = np.random.randn(500, 500)
并检查性能:In [77]: %timeit sumsqdiff(z, template)
1 loops, best of 3: 3.55 s per loop
In [78]: %timeit sumsqdiff2(z, template)
10 loops, best of 3: 33 ms per loop
不太寒酸。:)
两个缺点:sumsqdiff2中的计算将生成一个临时数组,对于3x3模板,该数组的大小是input_image的9倍。(通常它是template.size倍于input_image的大小。)
当你对代码进行计算机化处理时,这些“跨步技巧”不会对你有所帮助。当转换为Cython时,您通常会回到使用numpy进行矢量化时消除的循环。在