def kp_selection_good_depth(self, cur_data, ref_data,
rigid_kp_score_method):
"""Choose valid kp from a series of operations
Args:
cur_data (dict): current data
ref_data (dict): reference data
rigid_kp_score_method (str): [opt_flow, rigid_flow]
Returns:
a dictionary containing
- **kp1_depth** (array, [Nx2]): keypoints in view-1
- **kp2_depth** (array, [Nx2]): keypoints in view-2
- **rigid_flow_mask** (array, [HxW]): rigid-optical flow consistency
"""
outputs = {}
# initialization
h, w = cur_data['depth'].shape
kp1 = image_grid(h, w)
kp1 = np.expand_dims(kp1, 0)
tmp_flow_data = np.transpose(np.expand_dims(ref_data['flow'], 0),
(0, 2, 3, 1))
kp2 = kp1 + tmp_flow_data
""" opt-rigid flow consistent kp selection """
if self.cfg.kp_selection.rigid_flow_kp.enable:
# compute rigid flow
rigid_flow_pose = ref_data['rigid_flow_pose'].pose
# Compute rigid flow
pose_tensor = torch.from_numpy(rigid_flow_pose).float().unsqueeze(
0).cuda()
depth = torch.from_numpy(ref_data['raw_depth']).float().unsqueeze(
0).unsqueeze(0).cuda()
rigid_flow_tensor = self.rigid_flow_layer(
depth,
pose_tensor,
self.K,
self.inv_K,
normalized=False,
)
rigid_flow = rigid_flow_tensor.detach().cpu().numpy()[0]
# compute optical-rigid flow difference
rigid_flow_diff = np.linalg.norm(rigid_flow - ref_data['flow'],
axis=0)
ref_data['rigid_flow_diff'] = np.expand_dims(rigid_flow_diff, 2)
# get depth-flow consistent kp
outputs.update(
opt_rigid_flow_kp(kp1=kp1,
kp2=kp2,
ref_data=ref_data,
cfg=self.cfg,
outputs=outputs,
score_method=rigid_kp_score_method))
return outputs
def kp_selection(self, cur_data, ref_data):
"""Choose valid kp from a series of operations
Args:
cur_data (dict): data of current frame (view-2)
ref_data (dict): data of reference frame (view-1)
Returns:
outputs (dict): a dictionary containing some of the following items
- **kp1_best** (array, [Nx2]): keypoints on view-1
- **kp2_best** (array, [Nx2]): keypoints on view-2
- **kp1_list** (array, [Nx2]): keypoints on view-1
- **kp2_list** (array, [Nx2]): keypoints on view-2
- **kp1_depth** (array, [Nx2]): keypoints in view-1
- **kp2_depth** (array, [Nx2]): keypoints in view-2
- **rigid_flow_mask** (array, [HxW]): rigid-optical flow consistency
"""
outputs = {}
outputs['good_kp_found'] = True
# initialization
h, w = cur_data['depth'].shape
kp1 = image_grid(h, w)
kp1 = np.expand_dims(kp1, 0)
tmp_flow_data = np.transpose(np.expand_dims(ref_data['flow'], 0), (0, 2, 3, 1))
kp2 = kp1 + tmp_flow_data
""" best-N selection """
if self.cfg.kp_selection.local_bestN.enable:
kp_sel_method = local_bestN
outputs.update(
kp_sel_method(
kp1=kp1,
kp2=kp2,
ref_data=ref_data,
cfg=self.cfg,
outputs=outputs
)
)
elif self.cfg.kp_selection.bestN.enable:
kp_sel_method = bestN_flow_kp
outputs.update(
kp_sel_method(
kp1=kp1,
kp2=kp2,
ref_data=ref_data,
cfg=self.cfg,
outputs=outputs
)
)
""" sampled kp selection """
if self.cfg.kp_selection.sampled_kp.enable:
outputs.update(
sampled_kp(
kp1=kp1,
kp2=kp2,
ref_data=ref_data,
kp_list=self.kps['uniform'],
cfg=self.cfg,
outputs=outputs
)
)
return outputs