python中生成面向光流直方图(HOOF)的最有效方法

Question

我有一个498帧长的图像序列，我使用cv2.calcOpticalFlowFarneback为其计算光流。因此，现在我有497个矢量图来表示我的运动矢量，这些矢量由大小和方向来描述。

我需要做的是生成一个直方图，在x轴上我有以度为单位的角度范围。更具体地说，我有12个bin，其中第一个bin包含方向为0 < angle < 30的所有向量，第二个bin包含30 < angle < 60的所有向量等等。相反，在y轴上，我需要得到每个bin中包含的这些向量的模的和。

这里的问题是，使用简单的for循环和if语句来完成所有这些需要很长时间：

#magnitude and angle are two np.array of shape (497, 506, 1378)
bins = [1,2,3,4,5,6,7,8,9,10,11,12]
sum = np.zeros_like(bins)
for idx in range(np.array(magnitude).shape[0]): # for each flow map, i.e. for each image pair
for mag, ang in zip(magnitude[idx].reshape(-1), angle[idx].reshape(-1)): 
if ang >= 0 and ang <= 30:
sum[0] += mag
elif ang > 30 and ang <= 60:
sum[1] += mag
elif ang > 60 and ang <= 90:
sum[2] += mag
elif ang > 90 and ang <= 120:
sum[3] += mag
elif ang > 120 and ang <= 150:
sum[4] += mag
elif ang > 150 and ang <= 180:
sum[5] += mag
elif ang > 180 and ang <= 210:
sum[6] += mag
elif ang > 210 and ang <= 240:
sum[7] += mag
elif ang > 240 and ang <= 270:
sum[8] += mag
elif ang > 270 and ang <= 300:
sum[9] += mag
elif ang > 300 and ang <= 330:
sum[10] += mag
elif ang > 330 and ang <= 360:
sum[11] += mag

计算大约需要3个小时。有人能提出一种更好、更有效的方法来进行这种计算吗？

提前谢谢。

---

编辑

去掉了条件句，并使用Numba来进一步加快速度。以下代码的计算时间不到10秒：

import numpy as np
from numba import jit
@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def hoof(magnitude, angle):
sum = np.zeros(13)
for idx in range(magnitude.shape[0]): # for each flow map, i.e. for each image pair
for mag, ang in zip(magnitude[idx].reshape(-1), angle[idx].reshape(-1)): 
sum[int((ang)//30)] += mag

sum[11] += sum[12]
return sum[0:12]

Answer 1

条件是缓慢的。你应该尽量避开他们。此外，Numpy矢量化和Numba JIT有助于大大加快此类代码的速度。下面是一个未经测试的例子：

import numba as nb
@nb.jit
def compute(magnitude, angle):
s = np.zeros(12)
for idx in range(magnitude.shape[0]):
for mag, ang in zip(magnitude[idx].reshape(-1), angle[idx].reshape(-1)):
if ang == 0:
s[0] += mag
elif ang > 0 and ang <= 360: # The condition can be removed if always true
s[(ang-1)//30] += mag
return s
# Assume both are Numpy array and angle is of type int.
# Note that the first call will be slower unless you precise the types.
compute(magnitude, angle)