def forward(self, input):
N = input.shape[0]
assert input.dtype == torch.float32
assert input.shape[-3:] == self.input_shape_chw
assert N == self.aggr_pyramid_f_a[0].shape[0]
# Aggregate pyramid feature on frame B
aggr_pyramid_f_b = self.aggregate_pyramid_features(
self.backbone_net.forward(input))
T = torch.eye(4).detach()
for level in [2, 1, 0]:
(level_H, level_W) = self.level_dim_hw[level]
M = self.sel_a_indices.shape[2] # number of selected pts
# Features on current level
f_a = self.aggr_pyramid_f_a[level]
f_b = aggr_pyramid_f_b[level]
f_b_grad = batched_gradient(f_b) # dim: (N, 2*C, H, W)
# Resize and Rescale the depth and the intrinsic matrix
rescale_ratio = 1.0 / math.pow(2, level)
level_K = rescale_ratio * self.K.detach() # dim: (N, 3, 3)
level_d_a = F.interpolate(
self.d_a,
scale_factor=rescale_ratio).detach() # dim: (N, 1, H, W)
sel_a_idx = self.sel_a_indices[:, level, :].view(
N, M).detach() # dim: (N, M)
# Cache several variables:
x_a_2d = self.x_valid_2d[level] # dim: (N, H*W, 2)
X_a_3d = batched_pi_inv(level_K, x_a_2d,
level_d_a.view((N, level_H * level_W, 1)))
X_a_3d_sel = batched_index_select(X_a_3d, 1,
sel_a_idx) # dim: (N, M, 3)
# Run iteration 3 times
for itr in range(0, 1):
T, r, delta_norm, lamb = module.dm_levenberg_marquardt_itr(
T, X_a_3d, X_a_3d_sel, f_a, sel_a_idx, level_K, f_b,
f_b_grad, self.lambda_prediction, level)
print("Level:", level, " Itr", itr, " delta norm:", delta_norm)
return T[:, :3, :3], T[:, :3, 3]
def valid(self, I_a, d_a, sel_a_indices, K, I_b, se3_gt, epoch):
"""
Pre cache the variable for prediction
:param I_a: Image of frame A, dim: (N, C, H, W)
:param d_a: Depth of frame A, dim: (N, 1, H, W)
:param sel_a_indices: (N, 3, M)
:param K: intrinsic matrix at level 0: dim: (N, 3, 3)
:param I_b: Image of frame B, dim: (N, C, H, W)
:param se3_gt: ground truth Pose
"""
(N, C, H, W) = I_a.shape
I_a.detach()
I_b.detach()
# Ground-truth pose
R_gt, t_gt = se3_exp(se3_gt)
# Concate I_a and I_b
I = torch.cat([I_a, I_b], dim=0)
# Aggregate pyramid features
aggr_pyramid = self.aggregate_pyramid_features(
self.backbone_net.forward(I))
aggr_pyramid_f_a = [f[:N, :, :, :] for f in aggr_pyramid]
aggr_pyramid_f_b = [f[N:, :, :, :] for f in aggr_pyramid]
# Init a se(3) vector and mark requires_grad = True
# alpha = torch.tensor([1e-4, 1e-4, 1e-4, 0.0, 0.0, 0.0]).repeat(N).view((N, 6)) # dim: (N, 6)
# factor = 0.3
# alpha = module.gen_random_alpha(se3_gt, rot_angle_rfactor=1.25, trans_vec_rfactor=0.16).view((N, 6)).cuda()
# alpha.requires_grad_()
T = torch.eye(4).view(1, 4, 4).repeat(N, 1, 1).detach()
init_T = T
pred_SE3_list = [
] # (num_level: low_res to high_res, num_iter_per_level)
gt_f_pair_list = []
lambda_weight = []
flow_list = []
for level in [2, 1, 0]:
pred_SE3_list.append([])
lambda_weight.append([])
flow_list.append([])
(level_H, level_W) = self.level_dim_hw[level]
M = sel_a_indices.shape[2] # number of selected pts
# Features on current level
f_a = aggr_pyramid_f_a[level]
f_b = aggr_pyramid_f_b[level]
f_b_grad = batched_gradient(f_b) # dim: (N, 2*C, H, W)
# Resize and Rescale the depth and the intrinsic matrix
rescale_ratio = 1.0 / math.pow(2, level)
level_K = rescale_ratio * K.detach() # dim: (N, 3, 3)
level_d_a = F.interpolate(
d_a, scale_factor=rescale_ratio).detach() # dim: (N, 1, H, W)
sel_a_idx = sel_a_indices[:,
level, :].view(N,
M).detach() # dim: (N, M)
# Cache several variables:
x_a_2d = self.x_train_2d[level] # dim: (N, H*W, 2)
X_a_3d = batched_pi_inv(level_K, x_a_2d,
level_d_a.view((N, level_H * level_W, 1)))
X_a_3d_sel = batched_index_select(X_a_3d, 1,
sel_a_idx) # dim: (N, M, 3)
""" Ground-truth correspondence for Regularizer
"""
f_C = f_a.shape[1]
X_b_3d_gt = batched_transpose(R_gt, t_gt, X_a_3d)
x_b_2d_gt, _ = batched_pi(level_K, X_b_3d_gt)
x_b_2d_gt = module.batched_x_2d_normalize(float(level_H),
float(level_W),
x_b_2d_gt).view(
N, level_H, level_W,
2) # (N, H, W, 2)
gt_f_wrap_b = batched_interp2d(f_b, x_b_2d_gt)
f_a_select = batched_index_select(
f_a.view(N, f_C, level_H * level_W), 2, sel_a_idx)
gt_f_wrap_b_select = batched_index_select(
gt_f_wrap_b.view(N, f_C, level_H * level_W), 2, sel_a_idx)
gt_f_pair_list.append((f_a_select, gt_f_wrap_b_select))
# Run iteration 3 times
for itr in range(0, 6):
T, r, delta_norm, lamb, flow = module.dm_levenberg_marquardt_itr(
T, X_a_3d, X_a_3d_sel, f_a, sel_a_idx, level_K, f_b,
f_b_grad, self.lambda_prediction, level)
pred_SE3_list[-1].append(T)
flow_list[-1].append((flow, x_b_2d_gt.detach()))
return pred_SE3_list, gt_f_pair_list, init_T.detach(), flow_list