├── LK.py
└── README.md


/LK.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from torch.autograd import Variable as V
 3 | from torch import Tensor as T
 4 | import numpy as np
 5 | import torch.nn.functional as F
 6 | 
 7 | def VT(x): return V(T(x), requires_grad=False)
 8 | 
 9 | def three_conv(dx,dy,dz,dt, fac=1):
10 |     #Factor for adding that single minus on the dt conv
11 |     conv = torch.nn.Conv3d(1,4,2)
12 |     conv.weight = torch.nn.Parameter(T(np.concatenate([dx,dy,dz,fac * dt], axis=0)))
13 |     conv.bias = torch.nn.Parameter(T(np.array([0,0,0,0])))
14 |     return conv
15 | 
16 | def img_derivatives(img1, img2):
17 |     ones = np.ones((2,2,2))
18 |     dx = (0.25 * ones * np.array([-1, 1]))[None,None,...]
19 |     dy = (0.25 * ones * np.array([-1, 1])[:, None])[None,None,...]
20 |     dz = 0.25 * np.stack([-np.ones((2,2)), np.ones((2,2))])[None,None,...]
21 |     dt = ones[None, None,...]
22 |     
23 |     conv1 = three_conv(dx,dy,dz,dt)
24 |     conv2 = three_conv(dx,dy,dz,dt, fac=-1)
25 |     res =  0.5 * (conv1(VT(img1[None,...])) + conv2(VT(img2[None,...])))[0]
26 |     #Returns a 4,50,50,50 for the 4 derivatives including time
27 |     return F.pad(res, (1,0,1,0,1,0))
28 | 
29 | def opt_flow(dimg, r=2):
30 |     d = dimg.shape[-1]
31 |     x = np.ones((1,1,2,2,2))
32 |     calc = (dimg[None, 0:3,  ...] * dimg[:,None, ...])
33 |     conv_next = torch.nn.Conv3d(3,3,2)
34 |     conv_next.weight = torch.nn.Parameter(T(x))
35 |     conv_next.bias = torch.nn.Parameter(T(np.array([0])))
36 |     
37 |     sum_conv = torch.cat([conv_next(i[:,None,...]) for i in torch.unbind(calc, 1)], 1)
38 |     dim = sum_conv.shape[-1]
39 |     
40 |     a = sum_conv
41 |     b = a.permute(2, 3, 4, 0, 1)
42 |     c = b[..., :-1, :].contiguous().view(-1, 3, 3)
43 |     d = b[..., -1, :].contiguous().view(-1, 3, 1)
44 | 
45 |     inv = torch.stack([mat.inverse() for mat in torch.unbind(c, 0)])
46 |     out = torch.bmm(inv, d)
47 |     out = out.transpose(0,1).contiguous().view(3,dim,dim,dim)
48 |     return out


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Lucas Kanade 3D Optical Flow in Pytorch
2 | 
3 | Pytorch implementation of Lucas Kanade 3D optical flow method
4 | 


--------------------------------------------------------------------------------