def refine_with_predicted_bbox(pred,
observation,
intrinsics,
dist_coeffs,
gts=None,
threshold=5.,
ax=None
):
"""
Refine with predicted 3D cuboid (disabled by default).
"""
tempt_box_pred = pred.copy()
tempt_box_pred[1:, :] += tempt_box_pred[0, :].reshape(1, 3)
# use the predicted 3D bounding box size for refinement
refined_prediction = ltr.pnp_refine(tempt_box_pred, observation, intrinsics,
dist_coeffs)
# discard the results if the refined solution is to far away from the initial position
distance = refined_prediction[:, 0] - tempt_box_pred[0, :]
distance = np.sqrt(np.sum(distance**2))
if distance > threshold:
return False, None
else:
# plotting
if ax is not None:
vp.plot_lines(ax,
refined_prediction[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='g')
return True, refined_prediction
def refine_with_perfect_size(pred,
observation,
intrinsics,
dist_coeffs,
gts,
threshold=5.,
ax=None
):
"""
Use the gt 3D box size for refinement to show the performance gain with
size regression.
If there is a nearby ground truth bbox, use its size.
pred [9, 3] gts[N, 9, 3]
"""
pred_center = pred[0, :].reshape(1,3)
distance = np.sqrt(np.sum((gts[:, 0, :] - pred_center)**2, axis=1))
minimum_idx = np.where(distance == distance.min())[0][0]
if distance[minimum_idx] > threshold:
return False, None
else:
# First align the box with gt size with the predicted box, then refine
tempt_box_pred = pred.copy()
tempt_box_pred[1:, :] += tempt_box_pred[0, :].reshape(1, 3)
tempt_box_gt = gts[minimum_idx].copy()
tempt_box_gt[1:, :] += tempt_box_gt[0, :].reshape(1, 3)
pseudo_box = ltr.procrustes_transform(tempt_box_gt.T, tempt_box_pred.T)
refined_prediction = ltr.pnp_refine(pseudo_box.T, observation, intrinsics,
dist_coeffs)
if ax is not None:
vp.plot_lines(ax,
pseudo_box[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='y',
linestyle='-.')
vp.plot_lines(ax,
refined_prediction[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='b',
linestyle='-.')
return True, refined_prediction
def visualize_lifting_results(data,
prediction,
target=None,
sample_num=None,
intrinsics=None,
refine=False,
dist_coeffs=np.zeros((4,1)),
meta_data=None
):
# only take the coordinates
if data.shape[1] > 18:
data = data[:, :18]
# use the ground truth translation if provided in the meta_data
if 'roots' in meta_data:
target = np.hstack([meta_data['roots'], target])
prediction = np.hstack([meta_data['roots'], prediction])
sample_num = sample_num if sample_num else len(prediction)
chosen = np.random.choice(len(prediction), sample_num, replace=False)
if target is not None:
assert len(target) == len(prediction)
p3d_gt_sample = target[chosen].reshape(sample_num, -1, 3)
else:
p3d_gt_sample = None
p3d_pred_sample = prediction[chosen].reshape(sample_num, -1, 3)
data_sample = data[chosen].reshape(sample_num, -1, 2)
# vp.plot_comparison_relative(p3d_pred_sample[:9, 3:],
# p3d_gt_sample[:9, 3:])
ax = vp.plot_scene_3dbox(p3d_pred_sample, p3d_gt_sample)
if not refine:
return
# refine 3D point prediction by minimizing re-projection errors
assert intrinsics is not None
for idx in range(sample_num):
prediction = p3d_pred_sample[idx]
tempt_box_pred = prediction.copy()
tempt_box_pred[1:, :] += tempt_box_pred[0, :].reshape(1, 3)
observation = data_sample[idx]
# use the predicted 3D bounding box size for refinement
refined_prediction = pnp_refine(tempt_box_pred, observation, intrinsics,
dist_coeffs)
vp.plot_lines(ax,
refined_prediction[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='g'
)
# use the gt 3D box size for refinement
# first align a box with gt size with the predicted box, then refine
if target is None:
continue
tempt_box_gt = p3d_gt_sample[idx].copy()
tempt_box_gt[1:, :] += tempt_box_gt[0, :].reshape(1, 3)
pseudo_box = procrustes_transform(tempt_box_gt.T, tempt_box_pred.T)
refined_prediction2 = pnp_refine(pseudo_box.T, observation, intrinsics,
dist_coeffs)
vp.plot_lines(ax,
pseudo_box[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='y'
)
vp.plot_lines(ax,
refined_prediction2[:, 1:].T,
vp.plot_3d_bbox.connections,
dimension=3,
c='b'
)
return