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()

X, Y = np.meshgrid(np.arange(ref.shape[1]), np.arange(ref.shape[0]))

imshow(ref, cmap=cm.gray, clim=(0,1))

step = 8

quiver(X[::step,::step], Y[::step,::step],
       vxs[::step,::step], vys[::step,::step],
       color='g', angles='xy', scale_units='xy', scale=0.25)

title('Estimated velocity field (x4 scale)')