pytorch-photometric-image-alignment-slow.py

import math
import numpy as np
from numpy.linalg import norm,inv,pinv
import cv2
from PIL import Image

depth = Image.load('depth.png')

# cam2world: intrinsic parameters
w = int(640)
h = int(480)
u_0 = w/2
v_0 = h/2
vfovd = 55
f = v_0/math.tan((vfovd/2)*3.1415926/180.0)

cx = u_0
cy = v_0

K = np.array([[f,0,cx],[0,f,cy],[0,0,1]])
iK = inv(K)

# Meshgrid of starting image pixel coordinates
u = np.linspace(0, w-1, w)
v = np.linspace(0, h-1, h)
uv, vv = np.meshgrid(u, v)

img_coords = np.transpose(np.dstack((uv,vv,np.ones((h,w)))),(2,0,1))
                       
img_coords_flat = img_coords.reshape((3,-1))

# Back-projection
x_hat_flat = np.dot(iK,img_coords_flat)

x_hat = x_hat_flat
norms = norm(x_hat,axis=0)

Xt = np.vstack((x_hat * 1/norms * depth.reshape(1,h*w), np.ones((1,h*w))))

import torch
from torch.autograd import Variable

x_ = Xt[:3,:]
p_ = np.zeros((6,1))
p = Variable(torch.Tensor(p_), requires_grad=True)
x = Variable(torch.Tensor(x_))
K = Variable(torch.Tensor(np.array([[f,0,cx],[0,f,cy],[0,0,1]])))
theta = torch.norm(p[0:3])
r = p[0:3] / theta.expand_as(p[0:3])
cosTheta = torch.cos(theta).expand_as(torch.eye(3))
oneMinuscosTheta = (1-torch.cos(theta)).expand_as(torch.eye(3))
sinTheta = torch.sin(theta).expand_as(torch.eye(3))
rx = r[0].expand_as(torch.eye(3))
ry = r[1].expand_as(torch.eye(3))
rz = r[2].expand_as(torch.eye(3))
    
Rx = Variable(torch.Tensor([[0,0,0],[0,0,-1],[0,1,0]]))*rx
Ry = Variable(torch.Tensor([[0,0,1],[0,0,0],[-1,0,0]]))*ry
Rz = Variable(torch.Tensor([[0,-1,0],[1,0,0],[0,0,0]]))*rz
    
R1 = cosTheta * Variable(torch.eye(3))
R2 = oneMinuscosTheta*torch.ger(r.squeeze(),r.squeeze())
R3 = (Rx+Ry+Rz)*sinTheta
    
R = R1+R2+R3
    
x2 = torch.mm(R,x) + p[3:6].expand_as(x)

x3 = torch.mm(K,x2)
piv = x3[1]/x3[2]
piu = x3[0]/x3[2]
P = np.vstack((piu,piv))

img = Image.load('img.png')
# define accumulator
accum = np.zeros((h,w))

# cv2 remap
P_reshape = P.reshape(2,h,w)
Py = P_reshape[1].astype(np.float32)
Px = P_reshape[0].astype(np.float32)
accum = cv2.remap(img,Px,Py,cv2.INTER_LINEAR)

pgrad = []
for pi in [piu,piv]:
    m = x_.shape[1]
    Jacobian = torch.Tensor(m,6).zero_()
    for n in range(m):
        grad_mask = torch.zeros(m)
        grad_mask[n] = 1
        pi.backward(grad_mask, retain_variables=True)
        Jacobian[n,:] = p.grad.data
        p.grad.data.zero_()
    pgrad.append(Jacobian)

# Calculate image gradient
graduimg = cv2.Sobel(img,cv2.CV_32F,1,0,ksize=int(31))
gradvimg = cv2.Sobel(img,cv2.CV_32F,0,1,ksize=int(31))    
wGradu_flat = gradximg.reshape((1,-1))
wGradv_flat = gradyimg.reshape((1,-1))

# Prepare vectorizing operations
pgrad_u = pgrad[0]
pgrad_v = pgrad[1]
Wp = np.vstack([pgrad_u.numpy(),pgrad_v.numpy()])
delI = np.hstack([np.diag(wGradu_flat),np.diag(wGradv_flat)])

# get Jacobian
Jac = delI.dot(Wp)

# solve for parameter update
# residual error
r = (accum - img).reshape(-1,1)
delP = pinv(Jac).dot(r)
p_ = p_ - 1*delP