from pylab import *
import datasets
ref, src = datasets.regframes('traffic')
import dtcwt
transform = dtcwt.Transform2d()
ref_t = transform.forward(ref, nlevels=6)
src_t = transform.forward(src, nlevels=6)
import dtcwt.registration as registration

reg = registration.estimatereg(src_t, ref_t)
warped_src = registration.warp(src, reg, method='bilinear')

vxs, vys = registration.velocityfield(reg, ref.shape[:2], method='bilinear')
vxs = vxs * ref.shape[1]
vys = vys * ref.shape[0]

figure()
imshow(np.abs(vxs + 1j*vys), cmap=cm.hot)
title('Velocity field magnitude')