104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
|
|
""""Generate stereo baselines for kitti evaluation"""
|
|
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
|
|
from ..utils import get_keypoints, mask_joint_disparity, disparity_to_depth
|
|
|
|
|
|
def baselines_association(baselines, zzs, keypoints, keypoints_right, reid_features):
|
|
"""compute stereo depth for each of the given stereo baselines"""
|
|
|
|
# Initialize variables
|
|
zzs_stereo = defaultdict()
|
|
cnt_stereo = defaultdict(int)
|
|
|
|
features, features_r, keypoints, keypoints_r = factory_features(
|
|
keypoints, keypoints_right, baselines, reid_features)
|
|
|
|
# count maximum possible associations
|
|
cnt_stereo['max'] = min(keypoints.shape[0], keypoints_r.shape[0]) # pylint: disable=E1136
|
|
|
|
# Filter joints disparity and calculate avg disparity
|
|
avg_disparities, disparities_x, disparities_y = mask_joint_disparity(keypoints, keypoints_r)
|
|
|
|
# Iterate over each left pose
|
|
for key in baselines:
|
|
|
|
# Extract features of the baseline
|
|
similarity = features_similarity(features[key], features_r[key], key, avg_disparities, zzs)
|
|
|
|
# Compute the association based on features minimization and calculate depth
|
|
zzs_stereo[key] = np.empty((keypoints.shape[0]))
|
|
|
|
indices_stereo = [] # keep track of indices
|
|
best = np.nanmin(similarity)
|
|
while not np.isnan(best):
|
|
idx, arg_best = np.unravel_index(np.nanargmin(similarity), similarity.shape) # pylint: disable=W0632
|
|
zz_stereo, flag = disparity_to_depth(avg_disparities[idx, arg_best])
|
|
zz_mono = zzs[idx]
|
|
similarity[idx, :] = np.nan
|
|
indices_stereo.append(idx)
|
|
|
|
# Filter stereo depth
|
|
# if flag and verify_stereo(zz_stereo, zz_mono, disparities_x[idx, arg_best], disparities_y[idx, arg_best]):
|
|
if flag and (1 < zz_stereo < 80): # Do not add hand-crafted verifications to stereo baselines
|
|
zzs_stereo[key][idx] = zz_stereo
|
|
cnt_stereo[key] += 1
|
|
similarity[:, arg_best] = np.nan
|
|
else:
|
|
zzs_stereo[key][idx] = zz_mono
|
|
|
|
best = np.nanmin(similarity)
|
|
indices_mono = [idx for idx, _ in enumerate(zzs) if idx not in indices_stereo]
|
|
for idx in indices_mono:
|
|
zzs_stereo[key][idx] = zzs[idx]
|
|
zzs_stereo[key] = zzs_stereo[key].tolist()
|
|
|
|
return zzs_stereo, cnt_stereo
|
|
|
|
|
|
def factory_features(keypoints, keypoints_right, baselines, reid_features):
|
|
|
|
features = defaultdict()
|
|
features_r = defaultdict()
|
|
|
|
for key in baselines:
|
|
if key == 'reid':
|
|
features[key] = np.array(reid_features[0])
|
|
features_r[key] = np.array(reid_features[1])
|
|
else:
|
|
features[key] = np.array(keypoints)
|
|
features_r[key] = np.array(keypoints_right)
|
|
|
|
return features, features_r, np.array(keypoints), np.array(keypoints_right)
|
|
|
|
|
|
def features_similarity(features, features_r, key, avg_disparities, zzs):
|
|
|
|
similarity = np.empty((features.shape[0], features_r.shape[0]))
|
|
for idx, zz_mono in enumerate(zzs):
|
|
feature = features[idx]
|
|
|
|
if key == 'ml_stereo':
|
|
expected_disparity = 0.54 * 721. / zz_mono
|
|
sim_row = np.abs(expected_disparity - avg_disparities[idx])
|
|
|
|
elif key == 'pose':
|
|
# Zero-center the keypoints
|
|
uv_center = np.array(get_keypoints(feature, mode='center').reshape(-1, 1)) # (1, 2) --> (2, 1)
|
|
uv_centers_r = np.array(get_keypoints(features_r, mode='center').unsqueeze(-1)) # (m,2) --> (m, 2, 1)
|
|
feature_0 = feature[:2, :] - uv_center
|
|
feature_0 = feature_0.reshape(1, -1) # (1, 34)
|
|
features_r_0 = features_r[:, :2, :] - uv_centers_r
|
|
features_r_0 = features_r_0.reshape(features_r_0.shape[0], -1) # (m, 34)
|
|
sim_row = np.linalg.norm(feature_0 - features_r_0, axis=1)
|
|
|
|
else:
|
|
sim_row = np.linalg.norm(feature - features_r, axis=1)
|
|
|
|
similarity[idx] = sim_row
|
|
return similarity
|