def verify_features(self, I_a, d_a, K, I_b, se3_gt, x, y, title):
"""
Extract feature pyramids f_a, f_b of I_a and I_b
Wrap f_b to f_a
Compute distances of a pixel in f_a with the neighbors of its corresponding pixels in f_b
: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 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: Groundtruth of se3, dim: (N, 6)
:return:
"""
import banet_track.ba_debug as debug
(N, C, H, W) = I_a.shape
I_a.requires_grad_()
I_b.requires_grad_()
# 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]
for level in [2, 1, 0]:
(level_H, level_W) = self.level_dim_hw[level]
# 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)
# Cache several variables:
R, t = se3_exp(se3_gt)
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_b_3d = batched_transpose(R, t, X_a_3d)
x_b_2d, _ = batched_pi(level_K, X_b_3d)
x_b_2d = module.batched_x_2d_normalize(float(level_H), float(level_W), x_b_2d).view(N, level_H, level_W, 2) # (N, H, W, 2)
# Wrap the feature
level_aggr_pyramid_f_b_wrap = batched_interp2d(aggr_pyramid_f_b[level], x_b_2d)
level_x = int(x * rescale_ratio)
level_y = int(y * rescale_ratio)
left = level_x - debug.similar_window_offset
left = left if left >= 0 else 0
right = level_x + debug.similar_window_offset
up = level_y - debug.similar_window_offset
up = up if up >= 0 else 0
down = level_y + debug.similar_window_offset
batch_distance = torch.norm(aggr_pyramid_f_a[level][:, :, up:down, left:right] - # (N, level_H, level_W)
level_aggr_pyramid_f_b_wrap[:, :, level_y:level_y+1, level_x:level_x+1], 2, 1)
show_multiple_img([{'img': I_a[0].detach().cpu().numpy().transpose(1, 2, 0), 'title': 'I_a'},
{'img': I_b[0].detach().cpu().numpy().transpose(1, 2, 0), 'title': 'I_b'},
{'img': batch_distance[0].detach().cpu().numpy(), 'title': 'feature distance', 'cmap':'gray'}],
title=title, num_cols=3)
I_a = I_a.cuda()
d_a = d_a.cuda()
sel_a_indices = sel_a_indices.cuda()
K = K.cuda()
I_b = I_b.cuda()
q_gt = q_gt.cuda()
t_gt = t_gt.cuda()
se3_gt = se3_gt.cuda()
T_gt = T_gt.cuda()
# Forward
pred_pose, gt_f_pairs, init_T, lambda_weight, flow_list, gt_f_patch_pairs = net.train_feed(
I_a, d_a, sel_a_indices, K, I_b, se3_gt, epoch)
# Compute Losses
R_gt, t_gt = se3_exp(se3_gt)
rot_error, trans_error = pose_loss(pred_pose, q_gt, t_gt,
np.linspace(0.5, 1.3, num=9))
reg_loss = torch.tensor(0.0) #feature_reg_loss(gt_f_pairs)
flow, flow_gt = flow_list[-1][-1]
# mask out faraway wrapping points
flow_norm = torch.norm(flow - flow_gt, p=2, dim=3)
flow_mask = torch.lt(flow_norm, 2.828)
flow_accu = torch.mean(torch.masked_select(flow_norm, flow_mask))
total_loss = 0.5 * rot_error + trans_error #+ 0.1 * flow_accu
# Compute Accuracy
T = pred_pose[-1][-1]
R, t = 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
def compute_pose(I_a,
d_a,
sel_a_idx,
I_b,
K,
alpha,
T_gt,
opt_max_itr=100,
opt_eps=1e-5):
# Debug assert
assert sel_a_indices.dtype == torch.int64
assert I_a.dtype == torch.float32
assert I_b.dtype == torch.float32
assert d_a.dtype == torch.float32
# Dimension
N, C, H, W = I_a.shape
# Pre-processing
# sel_a_idx = select_gradient_pixels(I_a, d_a, threshold=50.0)[: 2000]
M = sel_a_idx.shape[1]
I_b_grad = batched_gradient(
I_b
) # dim: (N, 2*C, H, W), (N, 0:C, H, W) = dI/dx, (N, C:2C, H, W) = dI/dy
assert H == d_a.shape[2]
assert W == d_a.shape[3]
# se(3) vector init
lambda_w = 0.2 * torch.ones(N, 6)
d_a = d_a.view((N, H * W, 1))
# Points' 3D Position at Frame a
x_a_2d = x_2d_coords_torch(N, H, W).view(N, H * W, 2)
X_a_3d = batched_pi_inv(K, x_a_2d, d_a)
X_a_3d_sel = batched_index_select(X_a_3d, 1, sel_a_idx)
# groundtruth wrap
alpha_gt = torch.tensor([0.5, 0.5, 0.5, 0.0, 0.0, 0.0]).repeat(N).view(
(N, 6))
R_gt, _ = se3_exp(alpha_gt)
I = torch.eye(3).view(1, 3, 3).expand(N, 3, 3).cuda()
zeros = torch.zeros_like(T_gt[:, :, 3]).cuda()
random_t = torch.zeros_like(T_gt[:, :, 3]).normal_(std=0.001)
#print('random_t:', random_t)
X_b_3d_gt = batched_transpose(R_gt, zeros, X_a_3d_sel)
x_b_2d_gt, _ = batched_pi(K, X_b_3d_gt)
for itr in range(0, opt_max_itr):
R, t = se3_exp(alpha)
X_b_3d = batched_transpose(R, t, X_a_3d_sel)
x_b_2d, _ = batched_pi(K, X_b_3d)
# Residual error
e = (x_b_2d_gt - x_b_2d).view(N, M * 2) # (N, H*W*2)
# Compute Jacobin Mat.
# Jacobi of Camera Pose: delta_u / delta_alpha
J = -J_camera_pose(X_a_3d_sel, K).view(N, M * 2, 6) # (N*M, 2, 6)
# x_b_2d = batched_x_2d_normalize(H, W, x_b_2d).view(N, H, W, 2) # (N, H, W, 2)
#
# # Wrap the image
# I_b_wrap = batched_interp2d(I_b, x_b_2d)
#
# # Residual error
# e = (I_a - I_b_wrap).view(N, C, H*W) # (N, C, H, W)
# e = batched_index_select(e, 2, sel_a_idx) # (N, C, M)
# e = e.transpose(1, 2).contiguous().view(N, M*C) # (N, M, C)
#
# # Compute Jacobin Mat.
# # Jacobi of Camera Pose: delta_u / delta_alpha
# du_d_alpha = J_camera_pose(X_a_3d_sel, K).view(N * M, 2, 6) # (N*M, 2, 6)
#
# # Jacobi of Image gradient: delta_I_b / delta_u
# dI_du = batched_interp2d(I_b_grad, x_b_2d) # (N, 2*C, H, W)
# dI_du = batched_index_select(dI_du.view(N, 2*C, H*W), 2, sel_a_idx) # (N, 2*C, M)
# dI_du = torch.transpose(dI_du, 1, 2).contiguous().view(N * M, 2, C) # (N*M, 2, C)
# dI_du = torch.transpose(dI_du, 1, 2) # (N*M, C, 2)
#
# # J = -dI_b/du * du/d_alpha
# J = -torch.bmm(dI_du, du_d_alpha).view(N, C*M, 6)
# Compute the update parameters
delta, delta_norm = gauss_newtown_update(J, e) # (N, 6), (N, 1)
max_norm = torch.max(delta_norm).item()
if max_norm < opt_eps:
print('break')
break
r_norm = torch.sum(e * e, dim=1) / M #2.0
print('Itr:', itr, 'r_norm=', torch.sqrt(r_norm), "update_norm=",
max_norm)
# Update the delta
alpha = alpha + delta
return R, t
