def camera(transform: Isometry = Isometry(),
wh_ratio: float = 4.0 / 3.0,
scale: float = 1.0,
fovx: float = 90.0,
color_id: int = -1):
pw = np.tan(np.deg2rad(fovx / 2.)) * scale
ph = pw / wh_ratio
all_points = np.asarray([
[0.0, 0.0, 0.0],
[pw, ph, scale],
[pw, -ph, scale],
[-pw, ph, scale],
[-pw, -ph, scale],
])
line_indices = np.asarray([[0, 1], [0, 2], [0, 3], [0, 4], [1, 2], [1, 3],
[3, 4], [2, 4]])
geom = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(all_points),
lines=o3d.utility.Vector2iVector(line_indices))
if color_id == -1:
my_color = np.zeros((3, ))
else:
my_color = np.asarray(
matplotlib.cm.get_cmap('tab10').colors)[color_id, :3]
geom.colors = o3d.utility.Vector3dVector(
np.repeat(np.expand_dims(my_color, 0), line_indices.shape[0], 0))
geom.transform(transform.matrix)
return geom
def compute_sdf_Hg(self,
n_iter: int,
last_pose: Isometry,
cur_delta_pose: Isometry,
obs_xyz: torch.Tensor,
no_grad: bool = False):
"""
Get the error function: L(xi) = mu_[T(xi)p] / std_[T(xi)p].detach()
:param cur_pose:
:param obs_xyz: (N, 3) in camera space.
:return: f: (M, ) J: (M, 6)
"""
cur_obs_xyz = (last_pose.dot(cur_delta_pose)) @ obs_xyz
cur_obs_xyz.requires_grad_(not no_grad)
cur_obs_sdf, cur_obs_std, cur_obs_valid_mask = self.map.get_sdf(
cur_obs_xyz)
# print(cur_obs_std.min(), cur_obs_std.max()) # ~0.13
cur_obs_sdf = cur_obs_sdf / cur_obs_std.detach()
if no_grad:
JW = None
else:
dsdf_dpos = torch.autograd.grad(
cur_obs_sdf, [cur_obs_xyz],
grad_outputs=torch.ones_like(cur_obs_sdf),
retain_graph=False,
create_graph=False)[0] # (N, 3)
del cur_obs_xyz
dsdf_dpos = dsdf_dpos[cur_obs_valid_mask] # (M, 3)
cur_obs_sdf = cur_obs_sdf.detach()
cur_dxyz = (cur_delta_pose @ obs_xyz)[cur_obs_valid_mask]
Lt = torch.from_numpy(
last_pose.q.rotation_matrix.astype(np.float32).T).cuda()
Lai = torch.mm(dsdf_dpos, Lt)
Lbi = torch.cross(cur_dxyz, Lai)
dsdf_dxi = torch.cat([Lai, Lbi], dim=-1)
JW = dsdf_dxi
Wf = cur_obs_sdf
# print("[Empirical] Robust kernel should be = ", 2.6 * torch.std(Wf))
if self.sdf_args.robust_kernel is not None:
term_weight = self._robust_kernel(cur_obs_sdf,
self.sdf_args.robust_kernel,
self.sdf_args.robust_k)
Wf = Wf * term_weight # (M)
JW = JW * term_weight.unsqueeze(1) if JW is not None else None
error_scale = 1.0 / Wf.size(0)
sum_error = (cur_obs_sdf * Wf).sum().item() * error_scale
if no_grad:
return None, None, float(sum_error)
else:
H = torch.einsum('na,nb->nab', JW, dsdf_dxi).sum(0) * error_scale
g = (dsdf_dxi * Wf.unsqueeze(1)).sum(0) * error_scale
return H.cpu().numpy().astype(float), g.cpu().numpy().astype(
float), float(sum_error)
def compute_flow(base_pc: np.ndarray, base_segms: np.ndarray, base_cam: Isometry, base_motions: list,
dest_cam: Isometry, dest_motions: list):
n_parts = len(base_motions)
final_pc = np.empty_like(base_pc)
for part_id in range(n_parts):
part_mask = np.where(base_segms == (part_id + 1))[0]
part_pc = (dest_cam.inv().dot(dest_motions[part_id]).dot(
base_motions[part_id].inv()).dot(base_cam)) @ base_pc[part_mask]
final_pc[part_mask] = part_pc
return final_pc - base_pc
def __getitem__(self, data_id):
idx, view_sel_idx = data_id // len(self.view_sel), data_id % len(self.view_sel)
pcs, segms, trans_dict = self._get_item(idx)
n_parts = len(trans_dict) - 1
n_views = pcs.shape[0]
view_sel = self.view_sel[view_sel_idx]
if view_sel is None:
view_sel = list(range(n_views))
def get_view_motions(view_id):
return [Isometry.from_matrix(trans_dict[t][view_id]) for t in range(1, n_parts + 1)]
ret_vals = {}
if DatasetSpec.PC in self.spec:
ret_vals[DatasetSpec.PC] = pcs[view_sel, ...]
if DatasetSpec.SEGM in self.spec:
ret_vals[DatasetSpec.SEGM] = segms[view_sel, ...]
if DatasetSpec.FLOW in self.spec:
assert len(view_sel) == 2
view0_id, view1_id = view_sel
flow12 = compute_flow(pcs[view0_id], segms[view0_id],
Isometry.from_matrix(trans_dict['cam'][view0_id]), get_view_motions(view0_id),
Isometry.from_matrix(trans_dict['cam'][view1_id]), get_view_motions(view1_id))
ret_vals[DatasetSpec.FLOW] = flow12
if DatasetSpec.FULL_FLOW in self.spec:
all_flows = []
for view_i in view_sel:
for view_j in view_sel:
flow_ij = compute_flow(pcs[view_i], segms[view_i],
Isometry.from_matrix(trans_dict['cam'][view_i]), get_view_motions(view_i),
Isometry.from_matrix(trans_dict['cam'][view_j]), get_view_motions(view_j))
all_flows.append(flow_ij)
all_flows = np.stack(all_flows)
ret_vals[DatasetSpec.FULL_FLOW] = all_flows
return [ret_vals[k] for k in self.spec]
def track_camera_points_lm(self, init_pose: Isometry,
obs_xyz: torch.Tensor):
assert obs_xyz.size(1) == 3
cur_pose = init_pose
damping = self.args.lm_damping_init
for i_iter in range(self.args.n_gn_iter):
f, dsdf_dxi = self.get_error_func(cur_pose,
obs_xyz,
need_grad=True) # (M, ), (M, 6)
print(torch.std(f))
term_weight = torch.ones_like(f)
if self.args.robust_kernel:
obs_sdf_abs = torch.abs(f) # (M,)
term_weight = torch.where(obs_sdf_abs <= self.args.robust_k,
term_weight,
self.args.robust_k / obs_sdf_abs)
Wf = f * term_weight # (M, 1)
H = torch.einsum('nx,ny->xy', dsdf_dxi * term_weight.unsqueeze(1),
dsdf_dxi).cpu().numpy() # (6, 6)
lambda_DtD = damping * np.diag(np.diag(H))
g = -(dsdf_dxi * Wf.unsqueeze(1)).sum(0).cpu().numpy() # (6,)
xi = np.linalg.solve(H + lambda_DtD, g)
new_pose = Isometry.from_twist(xi) @ cur_pose
rho_denom = (xi * (lambda_DtD @ xi)).sum() + (xi * g).sum()
f_new = self.get_error_func(new_pose, obs_xyz,
need_grad=False) # (M, )
new_weight = torch.ones_like(f_new)
if self.args.robust_kernel:
f_abs = torch.abs(f_new)
new_weight = torch.where(f_abs <= self.args.robust_k,
new_weight,
self.args.robust_k / f_abs)
rho = ((f * Wf.squeeze()).sum() -
(f_new * new_weight * f_new).sum()).item() / rho_denom
if rho > self.args.lm_eps4:
damping /= self.args.lm_ldown
cur_pose = new_pose
print(f"LM Iter {i_iter} LM Accepted: {(f ** 2).sum().item()}")
else:
damping *= self.args.lm_lup
print(f"LM Iter {i_iter} LM Rejected: {(f ** 2).sum().item()}")
damping = min(max(damping, 1.0e-7), 1.0e7)
return cur_pose
def frame(transform: Isometry = Isometry(), size=1.0):
frame_obj = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size)
frame_obj.transform(transform.matrix)
return frame_obj
def track_camera(self,
rgb_data: torch.Tensor,
depth_data: torch.Tensor,
calib: FrameIntrinsic,
set_pose: Isometry = None):
"""
:param rgb_data: (H, W, 3) float32
:param depth_data: (H, W) float32
:param calib: FrameIntrinsic
:param set_pose: Force set pose.
:return: a pose.
"""
cur_intensity = torch.mean(rgb_data, dim=-1)
cur_depth = depth_data
cur_intensity, cur_depth, cur_dIdxy = self._make_image_pyramid(
cur_intensity, cur_depth)
cur_rgb = rgb_data.permute(2, 0, 1) # (3, H, W)
# Process to point cloud.
pc_scale = self.sdf_args.subsample
pc_data = torch.nn.functional.interpolate(
cur_depth[0].unsqueeze(0).unsqueeze(0),
scale_factor=pc_scale,
mode='nearest',
recompute_scale_factor=False).squeeze(0).squeeze(0)
cur_rgb = torch.nn.functional.interpolate(
cur_rgb.unsqueeze(0),
scale_factor=pc_scale,
mode='bilinear',
recompute_scale_factor=False).squeeze(0)
pc_data = unproject_depth(pc_data, calib.fx * pc_scale,
calib.fy * pc_scale, calib.cx * pc_scale,
calib.cy * pc_scale)
pc_data = torch.cat([
pc_data,
torch.zeros(
(pc_data.size(0), pc_data.size(1), 1), device=pc_data.device)
],
dim=-1)
pc_data = pc_data.reshape(-1, 4)
cur_rgb = cur_rgb.permute(1, 2, 0) # (W, H, 3)
cur_rgb = cur_rgb.reshape(-1, 3)
nan_mask = ~torch.isnan(pc_data[..., 0])
pc_data = pc_data[nan_mask]
cur_rgb = cur_rgb[nan_mask]
with torch.cuda.device(self.map.device):
pc_data_valid_mask = remove_radius_outlier(pc_data, 16, 0.05)
pc_data = pc_data[pc_data_valid_mask]
cur_rgb = cur_rgb[pc_data_valid_mask]
normal_data = estimate_normals(pc_data, 16, 0.1, [0.0, 0.0, 0.0])
normal_valid_mask = ~torch.isnan(normal_data[..., 0])
normal_data = normal_data[normal_valid_mask]
cur_rgb = cur_rgb[normal_valid_mask]
pc_data = pc_data[normal_valid_mask, :3]
self.last_colored_pcd = [pc_data, cur_rgb]
pc_data, normal_data = point_box_filter(pc_data, normal_data, 0.02)
self.last_processed_pc = [pc_data, normal_data]
if set_pose is not None:
final_pose = set_pose
else:
assert len(self.all_pd_pose) > 0
lspeed = Isometry()
final_pose = self.gauss_newton(self.all_pd_pose[-1].dot(lspeed),
cur_intensity, cur_depth, cur_dIdxy,
pc_data, calib)
self.last_intensity = cur_intensity
self.last_depth = cur_depth
self.all_pd_pose.append(final_pose)
return final_pose
def gauss_newton(self, init_pose: Isometry, cur_intensity_pyramid: list,
cur_depth_pyramid: list, cur_dIdxy_pyramid: list,
obs_xyz: torch.Tensor, calib: FrameIntrinsic):
last_pose = self.all_pd_pose[-1]
cur_delta_pose = last_pose.inv().dot(init_pose)
last_delta_pose = copy.deepcopy(cur_delta_pose)
i_iter = 0
for group_iter_config in self.args.iter_config:
last_energy = np.inf
from utils.exp_util import AverageMeter
loss_meter = AverageMeter()
for i_iter in list(range(group_iter_config["n"])) + [-1]:
H = np.zeros((6, 6), dtype=float)
g = np.zeros((6, ), dtype=float)
cur_energy = 0.0
for loss_config in group_iter_config["type"]:
if loss_config[0] == 'sdf':
sdf_H, sdf_g, sdf_energy = self.compute_sdf_Hg(
i_iter, last_pose, cur_delta_pose, obs_xyz,
i_iter == -1)
loss_meter.append_loss({'sdf': sdf_energy})
cur_energy += sdf_energy
if i_iter != -1:
H += sdf_H
g += sdf_g
if loss_config[0] == 'rgb':
pyramid_level = loss_config[1]
rgb_H, rgb_g, rgb_energy = self.compute_rgb_Hg(
pyramid_level, cur_delta_pose,
cur_intensity_pyramid, cur_depth_pyramid,
cur_dIdxy_pyramid, calib, i_iter == -1)
loss_meter.append_loss({'rgb': rgb_energy})
cur_energy += rgb_energy
if i_iter != -1:
H += rgb_H
g += rgb_g
if loss_config[0] == 'motion':
motion_H, motion_g, motion_energy = self.compute_motion_Hg(
cur_delta_pose, i_iter == -1)
loss_meter.append_loss({'motion': motion_energy})
cur_energy += motion_energy
if i_iter != -1:
H += motion_H
g += motion_g
if cur_energy > last_energy:
cur_delta_pose = last_delta_pose
break
else:
last_delta_pose = copy.deepcopy(cur_delta_pose)
last_energy = cur_energy
if i_iter != -1:
xi = np.linalg.solve(H, -g)
cur_delta_pose = Isometry.from_twist(xi) @ cur_delta_pose
# logging.info(f"GN Iter {i_iter} @ Group {group_iter_config}, Loss = {loss_meter.get_printable_newest()}")
if i_iter >= 10:
# Safe bar: more number of iterations indicate bad convergence. This may be due to too small regularization,
# So we fallback to our default setting after this detection.
self.n_unstable += 1
if self.n_unstable >= 3:
self.rgb_args.weight = max(self.rgb_args.weight, 500.)
return last_pose.dot(cur_delta_pose)
def get_view_motions(view_id):
return [Isometry.from_matrix(trans_dict[t][view_id]) for t in range(1, n_parts + 1)]