用颜色模型可视化光流

Question

我已经实现了一个密集光流算法，我想用下面的颜色模型来可视化它

（颜色表示某点的流向，强度表示位移矢量的长度）

我实现了一个肮脏的可视化版本

def visualizeFlow(u, v):
    colorModel = cv2.imread('../colormodel.png')
    colorModelCenter = (colorModel.shape[0]/2, colorModel.shape[1]/2)
    res = np.zeros((u.shape[0], u.shape[1], 3), dtype=np.uint8)
    mag = np.max(np.sqrt(u**2 + v**2)) 
    if mag == 0:
        return res, colorModel
    for i in xrange(res.shape[0]):
        for j in xrange(res.shape[1]):
            res[i, j] = colorModel[
                        colorModelCenter[0] + (v[i, j]/mag*colorModelCenter[0]),
                        colorModelCenter[1] + (u[i, j]/mag*colorModelCenter[1])
                    ]
    return res, colorModel

一般情况下它能拍出很好的照片，但速度很慢

所以我的问题是，有谁能帮我加快可视化的速度？如果有人知道一种更好的方法来可视化稠密的流动，那可能很酷